@marimo-team/frontend 0.23.13-dev0 → 0.23.13-dev2

package/src/core/codemirror/misc/paste.ts

@@ -1,6 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type { Extension } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
+import type { SyntaxNode } from "@lezer/common";
+import { parser } from "@lezer/python";
 import { cellActionsState } from "../cells/state";
 
 export function pasteBundle(): Extension[] {
@@ -8,152 +10,166 @@ export function pasteBundle(): Extension[] {
     EditorView.domEventHandlers({
       paste: (event: ClipboardEvent, view: EditorView) => {
         const text = event.clipboardData?.getData("text/plain");
-        if (!text?.includes("@app.cell")) {
+        if (!text || !looksLikeMarimoApp(text)) {
           return false;
         }
 
-        const cells = extractCells(text);
-        if (cells.length === 0) {
+        const { setup, cells } = extractMarimoApp(text);
+        if (setup === null && cells.length === 0) {
           return false;
         }
 
         const actions = view.state.facet(cellActionsState);
-        actions.createManyBelow(cells);
+        // The `with app.setup` block belongs in the special setup cell, not in
+        // an ordinary cell.
+        if (setup !== null) {
+          actions.addOrAppendSetupCell(setup);
+        }
+        if (cells.length > 0) {
+          actions.createManyBelow(cells);
+        }
         return true;
       },
     }),
   ];
 }
 
+/** Whether the pasted text looks like marimo app source worth parsing. */
+function looksLikeMarimoApp(text: string): boolean {
+  return /@app\.cell|@app\.function|\bapp\.setup\b/.test(text);
+}
+
+export interface ExtractedApp {
+  /** The `with app.setup` block's body, if present. Belongs in the setup cell. */
+  setup: string | null;
+  /** Ordinary cells, in source order. */
+  cells: string[];
+}
+
 /**
- * Extract the cells from a marimo app.
+ * Extract the cells of a marimo app from its source.
+ *
+ * Parses the pasted source with the Python grammar (the same one that powers
+ * the editor) and turns each marimo construct into a cell:
+ * - `with app.setup(...)` contributes the setup block's body, returned
+ *   separately as `setup` because it belongs in the dedicated setup cell.
+ * - `@app.cell` functions contribute their body, minus the trailing
+ *   auto-generated `return`.
+ * - `@app.function` definitions contribute the whole function (the decorator
+ *   is dropped, but the function — including its `return` — is kept).
+ *
+ * Using the real syntax tree means decorators, `async def`, multi-line
+ * signatures, nested functions, and multi-line returns of any bracket type are
+ * all handled structurally rather than by line-based heuristics.
  */
-export function extractCells(text: string): string[] {
-  // Quick check if this looks like a marimo app
-  if (!text.includes("@app.cell")) {
-    return [];
+export function extractMarimoApp(text: string): ExtractedApp {
+  if (!looksLikeMarimoApp(text)) {
+    return { setup: null, cells: [] };
   }
 
   const cells: string[] = [];
-  const lines = text.split("\n");
-  let currentCell: string[] = [];
-  let inCell = false;
-  let skipLines = 0;
-  let inMultilineArgs = false;
-  let inMultilineReturn = false;
-  let parenCount = 0;
-  let cellBaseIndent: number | null = null;
-
-  // Pre-compile regex patterns
-  const leadingParenRegex = /\(/g;
-  const trailingParenRegex = /\)/g;
-  const cellEndMarkers = new Set(["@"]);
-
-  function countParens(line: string): number {
-    return (
-      (line.match(leadingParenRegex) || []).length -
-      (line.match(trailingParenRegex) || []).length
-    );
-  }
-
-  function getIndent(line: string): number {
-    const match = line.match(/^\s*/);
-    return match ? match[0].length : 0;
-  }
-
-  function finalizeCellIfNeeded() {
-    if (currentCell.length === 0) {
-      return;
+  const setupBlocks: string[] = [];
+  const tree = parser.parse(text);
+
+  // A node's `from` starts after its line's leading indentation; extend back to
+  // the start of the line so `dedent` sees consistent indentation, then dedent.
+  const blockText = (from: number, to: number): string => {
+    const lineStart = text.lastIndexOf("\n", from - 1) + 1;
+    return dedent(text.slice(lineStart, to));
+  };
+
+  const pushCell = (from: number, to: number) => {
+    const cell = blockText(from, to);
+    if (cell.trim()) {
+      cells.push(cell);
     }
-
-    // Only add non-empty cells
-    if (currentCell.some((l) => l.trim() !== "")) {
-      cells.push(dedent(currentCell.join("\n")));
+  };
+
+  // The text of a block's body (the statements after its `:`), optionally
+  // dropping a trailing auto-generated `return`. Returns null if empty.
+  const bodyText = (
+    body: SyntaxNode,
+    stripTrailingReturn: boolean,
+  ): string | null => {
+    const children: SyntaxNode[] = [];
+    for (let child = body.firstChild; child; child = child.nextSibling) {
+      if (child.name !== ":") {
+        children.push(child);
+      }
     }
-    currentCell = [];
-  }
-
-  for (const line of lines) {
-    const trimmed = line.trim();
-
-    // Skip empty lines between cells
-    if (!trimmed && !inCell) {
-      continue;
+    if (stripTrailingReturn && children.at(-1)?.name === "ReturnStatement") {
+      children.pop();
     }
-
-    // Start of a new cell
-    if (trimmed.startsWith("@app.cell")) {
-      finalizeCellIfNeeded();
-      inCell = true;
-      skipLines = 1; // Skip the def line
-      cellBaseIndent = null;
-      continue;
+    const first = children[0];
+    const last = children.at(-1);
+    if (!first || !last) {
+      return null;
     }
-
-    // Handle function definition and args
-    if (skipLines > 0) {
-      if (
-        trimmed.startsWith("def") &&
-        trimmed.includes("(") &&
-        !trimmed.includes("):")
-      ) {
-        inMultilineArgs = true;
-        parenCount = countParens(trimmed);
+    const cell = blockText(first.from, last.to);
+    return cell.trim() ? cell : null;
+  };
+
+  tree.iterate({
+    enter: (node) => {
+      // Setup cell: `with app.setup:` / `with app.setup(...):`
+      if (node.name === "WithStatement") {
+        const body = node.node.getChild("Body");
+        if (body && text.slice(node.from, body.from).includes("app.setup")) {
+          const setup = bodyText(body, false);
+          if (setup !== null) {
+            setupBlocks.push(setup);
+          }
+        }
+        return false;
       }
-      skipLines--;
-      continue;
-    }
 
-    // Track multi-line arguments
-    if (inMultilineArgs) {
-      parenCount += countParens(trimmed);
-      if (parenCount === 0) {
-        inMultilineArgs = false;
+      if (node.name !== "DecoratedStatement") {
+        // Keep descending until we reach the decorated definitions.
+        return true;
       }
-      continue;
-    }
-
-    if (!inCell) {
-      continue;
-    }
-
-    // Handle cell content
-    if (cellEndMarkers.has(trimmed[0]) || trimmed.startsWith("if __name__")) {
-      finalizeCellIfNeeded();
-      inCell = trimmed.startsWith("@");
-      continue;
-    }
-
-    // Detect base indentation of cell body from first content line
-    if (cellBaseIndent === null && trimmed) {
-      cellBaseIndent = getIndent(line);
-    }
 
-    // Handle return statements — only strip cell-level returns
-    if (trimmed.startsWith("return") && getIndent(line) === cellBaseIndent) {
-      if (trimmed.includes("(") && !trimmed.endsWith(")")) {
-        inMultilineReturn = true;
-        parenCount = countParens(trimmed);
+      const decorator = node.node.getChild("Decorator");
+      const fn = node.node.getChild("FunctionDefinition");
+      if (!decorator || !fn) {
+        return false;
+      }
+      const decoratorText = text.slice(decorator.from, decorator.to).trim();
+
+      // `@app.cell` / `@app.cell(...)`: the body becomes a cell, minus the
+      // trailing auto-generated return.
+      if (decoratorText.startsWith("@app.cell")) {
+        const body = fn.getChild("Body");
+        if (body) {
+          const cell = bodyText(body, true);
+          if (cell !== null) {
+            cells.push(cell);
+          }
+        }
+        return false;
       }
-      continue;
-    }
 
-    if (inMultilineReturn) {
-      parenCount += countParens(trimmed);
-      if (parenCount === 0) {
-        inMultilineReturn = false;
+      // `@app.function`: the whole function definition is the cell (the
+      // decorator is dropped, but the function — including its return — stays).
+      if (decoratorText.startsWith("@app.function")) {
+        pushCell(fn.from, fn.to);
       }
-      continue;
-    }
 
-    // Add line to current cell
-    currentCell.push(line);
-  }
+      return false;
+    },
+  });
 
-  // Handle last cell
-  finalizeCellIfNeeded();
+  return {
+    setup: setupBlocks.length > 0 ? setupBlocks.join("\n\n") : null,
+    cells,
+  };
+}
 
-  return cells;
+/**
+ * Convenience wrapper returning only the ordinary cells (excluding the setup
+ * block). Prefer {@link extractMarimoApp} when the setup block matters.
+ */
+export function extractCells(text: string): string[] {
+  return extractMarimoApp(text).cells;
 }
 
 function dedent(text: string): string {