improvement: improve tab hanlding when in the middle of lines

frontend/src/core/codemirror/__tests__/tabs.test.ts

@@ -0,0 +1,177 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+import { describe, it, expect, vi, beforeAll, beforeEach } from "vitest";
+import { tabHandling, visibleForTesting } from "../tabs";
+import { EditorView } from "@codemirror/view";
+import { autocompletion, startCompletion } from "@codemirror/autocomplete";
+
+const { insertTab, startCompletionIfAtEndOfLine } = visibleForTesting;
+
+vi.mock("@codemirror/autocomplete", async (imported) => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const mod: any = await imported();
+  return {
+    ...mod,
+    startCompletion: vi.fn().mockImplementation(mod.startCompletion),
+  };
+});
+
+describe("insertTab", () => {
+  it("should insert 4 spaces when cursor is not in a selection", () => {
+    const doc = "Hello\nWorld";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+    });
+
+    const result = insertTab(view);
+
+    expect(result).toBe(true);
+    expect(view.state.doc.toString()).toBe("    Hello\nWorld");
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+
+  it("should call indentMore when there is a selection", () => {
+    const doc = "Hello\nWorld";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 0, head: 5 },
+    });
+
+    insertTab(view);
+
+    expect(view.state.doc.toString()).toBe("    Hello\nWorld");
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+});
+
+describe("startCompletionIfAtEndOfLine", () => {
+  beforeAll(() => {
+    // Mock getBoundingClientRect
+    // @ts-expect-error - JSDOM doesn't have createRange
+    document.createRange = vi.fn(() => ({
+      getBoundingClientRect: vi.fn(() => ({ width: 0 })),
+      setStart: vi.fn(),
+      setEnd: vi.fn(),
+      getClientRects: vi.fn(() => [{ width: 0 }]),
+    }));
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("should start completion at the end of a non-empty line", () => {
+    const doc = "Hello";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+
+      selection: { anchor: 5, head: 5 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+    expect(result).toBe(true);
+
+    expect(startCompletion).toHaveBeenCalled();
+  });
+
+  it("should not start completion if cursor is not at the end of a line", () => {
+    const doc = "Hello World";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 5, head: 5 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+
+    expect(result).toBe(false);
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+
+  it("should not start completion if line is empty", () => {
+    const doc = "\n";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 0, head: 0 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+
+    expect(result).toBe(false);
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+
+  it("should not start completion if cursor is in a selection", () => {
+    const doc = "Hello";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 0, head: 5 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+
+    expect(result).toBe(false);
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+
+  it("should not start completion if cursor is at the end of an empty line", () => {
+    const doc = "\n";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 1, head: 1 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+
+    expect(result).toBe(false);
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+
+  it("should not start completion if cursor is at the end of a line with only whitespace", () => {
+    const doc = "    ";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 4, head: 4 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+
+    expect(result).toBe(false);
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+
+  it("should not start completion if cursor is at the beginning of a line", () => {
+    const doc = "Hello\nWorld";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 6, head: 6 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+
+    expect(result).toBe(false);
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+
+  it("should not start completion if cursor is in whitespace", () => {
+    const doc = "Hello\n    World";
+    const view = new EditorView({
+      doc,
+      extensions: [autocompletion({}), tabHandling()],
+      selection: { anchor: 6, head: 6 },
+    });
+
+    const result = startCompletionIfAtEndOfLine(view);
+
+    expect(result).toBe(false);
+    expect(startCompletion).not.toHaveBeenCalled();
+  });
+});

frontend/src/core/codemirror/cm.ts

@@ -1,23 +1,12 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import {
-  acceptCompletion,
-  closeBrackets,
-  closeBracketsKeymap,
-  startCompletion,
-} from "@codemirror/autocomplete";
-import {
-  history,
-  historyKeymap,
-  indentWithTab,
-  indentMore,
-} from "@codemirror/commands";
+import { closeBrackets, closeBracketsKeymap } from "@codemirror/autocomplete";
+import { history, historyKeymap } from "@codemirror/commands";
 import {
   bracketMatching,
   defaultHighlightStyle,
   foldGutter,
   foldKeymap,
   indentOnInput,
-  indentUnit,
   syntaxHighlighting,
 } from "@codemirror/language";
 import {
@@ -56,6 +45,7 @@ import { historyCompartment } from "./editing/extensions";
 import { goToDefinitionBundle } from "./go-to-definition/extension";
 import type { HotkeyProvider } from "../hotkeys/hotkeys";
 import { lightTheme } from "./theme/light";
+import { tabHandling } from "./tabs";
 
 export interface CodeMirrorSetupOpts {
   cellId: CellId;
@@ -94,20 +84,6 @@ export const setupCodeMirror = (opts: CodeMirrorSetupOpts): Extension[] => {
   ];
 };
 
-const startCompletionAtEndOfLine = (cm: EditorView): boolean => {
-  const { from, to } = cm.state.selection.main;
-  if (from !== to) {
-    // this is a selection
-    return false;
-  }
-
-  const line = cm.state.doc.lineAt(to);
-  return line.text.slice(0, to - line.from).trim() === ""
-    ? // in the whitespace prefix of a line
-      false
-    : startCompletion(cm);
-};
-
 // Based on codemirror's basicSetup extension
 export const basicBundle = (opts: CodeMirrorSetupOpts): Extension[] => {
   const { theme, hotkeys, completionConfig } = opts;
@@ -141,7 +117,6 @@ export const basicBundle = (opts: CodeMirrorSetupOpts): Extension[] => {
     Prec.high(keymap.of(closeBracketsKeymap)),
     bracketMatching(),
     indentOnInput(),
-    indentUnit.of("    "),
     syntaxHighlighting(defaultHighlightStyle, { fallback: true }),
     keymap.of(foldKeymap),
 
@@ -152,20 +127,8 @@ export const basicBundle = (opts: CodeMirrorSetupOpts): Extension[] => {
     historyCompartment.of(history()),
     EditorState.allowMultipleSelections.of(true),
     findReplaceBundle(hotkeys),
-    keymap.of([
-      {
-        key: "Tab",
-        run: (cm) => {
-          return (
-            acceptCompletion(cm) ||
-            startCompletionAtEndOfLine(cm) ||
-            indentMore(cm)
-          );
-        },
-        preventDefault: true,
-      },
-    ]),
-    keymap.of([...historyKeymap, indentWithTab]),
+    tabHandling(),
+    keymap.of(historyKeymap),
   ];
 };
 

frontend/src/core/codemirror/tabs.ts

@@ -0,0 +1,81 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { acceptCompletion, startCompletion } from "@codemirror/autocomplete";
+import { indentMore, indentWithTab } from "@codemirror/commands";
+import { indentUnit } from "@codemirror/language";
+import type { Extension } from "@codemirror/state";
+import { keymap, type EditorView } from "@codemirror/view";
+
+/**
+ * Setup keybindings for tab handling.
+ */
+export function tabHandling(): Extension[] {
+  return [
+    indentUnit.of("    "),
+    keymap.of([
+      // Tabs for completion
+      {
+        key: "Tab",
+        run: (cm) => {
+          return (
+            acceptCompletion(cm) ||
+            startCompletionIfAtEndOfLine(cm) ||
+            insertTab(cm)
+          );
+        },
+        preventDefault: true,
+      },
+      // Tab-more/tab-less
+      indentWithTab,
+    ]),
+  ];
+}
+
+/**
+ * Custom insert tab that inserts our custom tab character (4 spaces).
+ * Adapted from https://github.com/codemirror/commands/blob/d0c97ba5de9d1b5d42fa5a713f0c9d64a3134a3c/src/commands.ts#L854-L858
+ */
+function insertTab(cm: EditorView) {
+  const { state } = cm;
+  if (state.selection.ranges.some((r) => !r.empty)) {
+    return indentMore(cm);
+  }
+  const indent = state.facet(indentUnit);
+  cm.dispatch(
+    state.update(state.replaceSelection(indent), {
+      scrollIntoView: true,
+      userEvent: "input",
+    }),
+  );
+  return true;
+}
+
+/**
+ * Start completion if the cursor is at the end of a line.
+ */
+function startCompletionIfAtEndOfLine(cm: EditorView): boolean {
+  const { from, to } = cm.state.selection.main;
+  if (from !== to) {
+    // this is a selection
+    return false;
+  }
+
+  const line = cm.state.doc.lineAt(to);
+  const textBeforeCursor = line.text.slice(0, to - line.from);
+
+  if (textBeforeCursor.trim() === "") {
+    // Cursor is at the beginning of a line or in whitespace
+    return false;
+  }
+
+  if (to === line.to) {
+    // Cursor is at the end of a line
+    return startCompletion(cm);
+  }
+
+  return false;
+}
+
+export const visibleForTesting = {
+  insertTab,
+  startCompletionIfAtEndOfLine,
+};