Add cell selection type (#246)

* Initial work

* Add support for "drink" / "pour" cell

* Add docstring

* Add jupyter dependency

* Add comments

src/actions/EditNewLine.ts

@@ -38,10 +38,18 @@ class EditNewLine implements Action {
     this.correctForParagraph(targets);
     if (this.isAbove) {
       await this.graph.actions.setSelectionBefore.run([targets]);
-      await commands.executeCommand("editor.action.insertLineBefore");
+      await commands.executeCommand(
+        targets[0].selectionContext.isNotebookCell
+          ? "jupyter.insertCellAbove"
+          : "editor.action.insertLineBefore"
+      );
     } else {
       await this.graph.actions.setSelectionAfter.run([targets]);
-      await commands.executeCommand("editor.action.insertLineAfter");
+      await commands.executeCommand(
+        targets[0].selectionContext.isNotebookCell
+          ? "jupyter.insertCellBelow"
+          : "editor.action.insertLineAfter"
+      );
     }
 
     return {

src/checkCommandValidity.ts

@@ -0,0 +1,26 @@
+import { ActionType, PartialTarget, SelectionType } from "./typings/Types";
+import { getPrimitiveTargets } from "./util/targetUtils";
+
+export function checkCommandValidity(
+  actionName: ActionType,
+  partialTargets: PartialTarget[],
+  extraArgs: any[]
+) {
+  if (
+    usesSelectionType("notebookCell", partialTargets) &&
+    !["editNewLineAbove", "editNewLineBelow"].includes(actionName)
+  ) {
+    throw new Error(
+      "The notebookCell scope type is currently only supported with the actions editNewLineAbove and editNewLineBelow"
+    );
+  }
+}
+
+function usesSelectionType(
+  selectionType: SelectionType,
+  partialTargets: PartialTarget[]
+) {
+  return getPrimitiveTargets(partialTargets).some(
+    (partialTarget) => partialTarget.selectionType === selectionType
+  );
+}

src/extension.ts

@@ -17,6 +17,7 @@ import { TestCase } from "./testUtil/TestCase";
 import { ThatMark } from "./core/ThatMark";
 import { TestCaseRecorder } from "./testUtil/TestCaseRecorder";
 import { getParseTreeApi } from "./util/getExtensionApi";
+import { checkCommandValidity } from "./checkCommandValidity";
 
 export async function activate(context: vscode.ExtensionContext) {
   const fontMeasurements = new FontMeasurements(context);
@@ -113,6 +114,8 @@ export async function activate(context: vscode.ExtensionContext) {
 
         const action = graph.actions[actionName];
 
+        checkCommandValidity(actionName, partialTargets, extraArgs);
+
         const targets = inferFullTargets(
           partialTargets,
           action.targetPreferences

src/processTargets/processSelectionType.ts

@@ -23,6 +23,8 @@ export default function (
   switch (target.selectionType) {
     case "token":
       return processToken(target, selection, selectionContext);
+    case "notebookCell":
+      return processNotebookCell(target, selection, selectionContext);
     case "document":
       return processDocument(target, selection, selectionContext);
     case "line":
@@ -32,6 +34,21 @@ export default function (
   }
 }
 
+function processNotebookCell(
+  target: PrimitiveTarget,
+  selection: SelectionWithEditor,
+  selectionContext: SelectionContext
+): TypedSelection {
+  const { selectionType, insideOutsideType, position } = target;
+  return {
+    selection,
+    selectionType,
+    position,
+    insideOutsideType,
+    selectionContext: { ...selectionContext, isNotebookCell: true },
+  };
+}
+
 function processToken(
   target: PrimitiveTarget,
   selection: SelectionWithEditor,

src/test/runTest.ts

@@ -7,6 +7,8 @@ import {
   downloadAndUnzipVSCode,
 } from "vscode-test";
 
+const extensionDependencies = ["pokey.parse-tree", "ms-toolsai.jupyter"];
+
 async function main() {
   try {
     // The folder containing the Extension Manifest package.json
@@ -22,9 +24,11 @@ async function main() {
       resolveCliPathFromVSCodeExecutablePath(vscodeExecutablePath);
 
     // Install extension dependencies
-    cp.spawnSync(cliPath, ["--install-extension", "pokey.parse-tree"], {
-      encoding: "utf-8",
-      stdio: "inherit",
+    extensionDependencies.forEach((dependency) => {
+      cp.spawnSync(cliPath, ["--install-extension", dependency], {
+        encoding: "utf-8",
+        stdio: "inherit",
+      });
     });
 
     // Run the integration test

src/test/suite/fixtures/recorded/selectionTypes/drinkCell.yml

@@ -0,0 +1,27 @@
+# Note that this is just checking that no errors are thrown
+spokenForm: drink cell
+languageId: python
+command:
+  actionName: editNewLineAbove
+  partialTargets:
+    - {type: primitive, selectionType: notebookCell}
+  extraArgs: []
+marks: {}
+initialState:
+  documentContents: |-
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 1, character: 12}
+      active: {line: 1, character: 12}
+finalState:
+  documentContents: |-
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 1, character: 0}
+      active: {line: 1, character: 0}
+  thatMark:
+    - anchor: {line: 1, character: 0}
+      active: {line: 1, character: 0}
+fullTargets: [{type: primitive, mark: {type: cursor}, selectionType: notebookCell, position: contents, insideOutsideType: inside, modifier: {type: identity}}]

src/test/suite/fixtures/recorded/selectionTypes/drinkCellEach.yml

@@ -0,0 +1,38 @@
+# Note that this is just checking that no errors are thrown
+spokenForm: drink cell each
+languageId: python
+command:
+  actionName: editNewLineAbove
+  partialTargets:
+    - type: primitive
+      selectionType: notebookCell
+      mark: {type: decoratedSymbol, symbolColor: default, character: e}
+  extraArgs: []
+marks:
+  default.e:
+    start: {line: 1, character: 7}
+    end: {line: 1, character: 12}
+initialState:
+  documentContents: |-
+    # %%
+    print("hello")
+
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 4, character: 12}
+      active: {line: 4, character: 12}
+finalState:
+  documentContents: |-
+    # %%
+    print("hello")
+
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 1, character: 0}
+      active: {line: 1, character: 0}
+  thatMark:
+    - anchor: {line: 1, character: 0}
+      active: {line: 1, character: 0}
+fullTargets: [{type: primitive, mark: {type: decoratedSymbol, symbolColor: default, character: e}, selectionType: notebookCell, position: contents, insideOutsideType: inside, modifier: {type: identity}}]

src/test/suite/fixtures/recorded/selectionTypes/pourCell.yml

@@ -0,0 +1,27 @@
+# Note that this is just checking that no errors are thrown
+spokenForm: pour cell
+languageId: python
+command:
+  actionName: editNewLineBelow
+  partialTargets:
+    - {type: primitive, selectionType: notebookCell}
+  extraArgs: []
+marks: {}
+initialState:
+  documentContents: |-
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 1, character: 12}
+      active: {line: 1, character: 12}
+finalState:
+  documentContents: |-
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 3, character: 0}
+      active: {line: 3, character: 0}
+  thatMark:
+    - anchor: {line: 3, character: 0}
+      active: {line: 3, character: 0}
+fullTargets: [{type: primitive, mark: {type: cursor}, selectionType: notebookCell, position: contents, insideOutsideType: inside, modifier: {type: identity}}]

src/test/suite/fixtures/recorded/selectionTypes/pourCellEach.yml

@@ -0,0 +1,38 @@
+# Note that this is just checking that no errors are thrown
+spokenForm: pour cell each
+languageId: python
+command:
+  actionName: editNewLineBelow
+  partialTargets:
+    - type: primitive
+      selectionType: notebookCell
+      mark: {type: decoratedSymbol, symbolColor: default, character: e}
+  extraArgs: []
+marks:
+  default.e:
+    start: {line: 1, character: 7}
+    end: {line: 1, character: 12}
+initialState:
+  documentContents: |-
+    # %%
+    print("hello")
+
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 4, character: 12}
+      active: {line: 4, character: 12}
+finalState:
+  documentContents: |-
+    # %%
+    print("hello")
+
+    # %%
+    print("hello")
+  selections:
+    - anchor: {line: 4, character: 0}
+      active: {line: 4, character: 0}
+  thatMark:
+    - anchor: {line: 4, character: 0}
+      active: {line: 4, character: 0}
+fullTargets: [{type: primitive, mark: {type: decoratedSymbol, symbolColor: default, character: e}, selectionType: notebookCell, position: contents, insideOutsideType: inside, modifier: {type: identity}}]

src/typings/Types.ts

@@ -58,7 +58,7 @@ export type Mark =
   | CursorMarkToken
   | That
   | Source
-//   | LastCursorPosition Not implemented yet
+  //   | LastCursorPosition Not implemented yet
   | DecoratedSymbol
   | LineNumber;
 export type Delimiter =
@@ -137,11 +137,8 @@ export type Modifier =
   | TailModifier;
 
 export type SelectionType =
-//   | "character" Not implemented
-  | "token"
-  | "line"
-  | "paragraph"
-  | "document";
+  //   | "character" Not implemented
+  "token" | "line" | "notebookCell" | "paragraph" | "document";
 export type Position = "before" | "after" | "contents";
 export type InsideOutsideType = "inside" | "outside" | null;
 
@@ -228,6 +225,8 @@ export interface SelectionContext {
    * The range of the delimiter after the selection
    */
   trailingDelimiterRange?: vscode.Range | null;
+
+  isNotebookCell?: boolean;
 }
 
 export interface TypedSelection {

src/util/targetUtils.ts

@@ -1,6 +1,10 @@
 import { TextEditor, Selection, Position } from "vscode";
 import { groupBy } from "./itertools";
-import { TypedSelection } from "../typings/Types";
+import {
+  PartialPrimitiveTarget,
+  PartialTarget,
+  TypedSelection,
+} from "../typings/Types";
 
 export function ensureSingleEditor(targets: TypedSelection[]) {
   if (targets.length === 0) {
@@ -86,3 +90,27 @@ function createTypeSelection(
     position: "contents",
   };
 }
+
+/**
+ * Given a list of targets, recursively descends all targets and returns every
+ * contained primitive target.
+ *
+ * @param targets The targets to extract from
+ * @returns A list of primitive targets
+ */
+export function getPrimitiveTargets(targets: PartialTarget[]) {
+  return targets.flatMap(getPrimitiveTargetsHelper);
+}
+
+function getPrimitiveTargetsHelper(
+  target: PartialTarget
+): PartialPrimitiveTarget[] {
+  switch (target.type) {
+    case "primitive":
+      return [target];
+    case "list":
+      return target.elements.flatMap(getPrimitiveTargetsHelper);
+    case "range":
+      return [target.start, target.end];
+  }
+}