Clicking filenames in tracebacks opens them in user's editor

Author: yairchu

frontend/src/components/editor/output/MarimoTracebackOutput.tsx

@@ -108,7 +108,7 @@ export const MarimoTracebackOutput = ({
             Fix with AI
           </Button>
         )}
-        {tracebackInfo && !isWasm() && (
+        {tracebackInfo && tracebackInfo.kind === "cell" && !isWasm() && (
           <Tooltip content={"Attach pdb to the exception point."}>
             <Button
               size="xs"
@@ -180,7 +180,7 @@ function lastLine(text: string): string {
 
 export const replaceTracebackFilenames = (domNode: DOMNode) => {
   const info = getTracebackInfo(domNode);
-  if (info) {
+  if (info?.kind === "cell") {
     const tooltipContent = <InsertBreakpointContent />;
     return (
       <span className="nb">
@@ -211,6 +211,18 @@ export const replaceTracebackFilenames = (domNode: DOMNode) => {
       </span>
     );
   }
+  if (info?.kind === "file") {
+    return (
+      <div
+        className="inline-block cursor-pointer text-destructive hover:underline"
+        onClick={(_) => {
+          getRequestClient().openFile({ path: info.filePath });
+        }}
+      >
+        <span className="nb">"{info.filePath}"</span>
+      </div>
+    );
+  }
 };
 
 export const replaceTracebackPrefix = (domNode: DOMNode) => {

frontend/src/core/codemirror/cells/traceback-decorations.ts

@@ -29,7 +29,7 @@ function createErrorDecorations(state: EditorState, errors: TracebackInfos) {
 
   // Filter and sort errors by line number to ensure they're added in order
   const relevantErrors = errors
-    .filter((error) => error.cellId === cellId)
+    .filter((error) => error.kind === "cell" && error.cellId === cellId)
     .sort((a, b) => a.lineNumber - b.lineNumber);
 
   for (const error of relevantErrors) {

frontend/src/utils/__tests__/traceback.test.ts

@@ -6,14 +6,20 @@ import { extractAllTracebackInfo } from "../traceback";
 describe("traceback", () => {
   test("extracts cell-link", () => {
     const errors = extractAllTracebackInfo(Tracebacks.raw);
-    expect(errors).toMatchInlineSnapshot(`
+    expect(
+      errors[0].kind === "file" &&
+        errors[0].filePath.endsWith("marimo/_runtime/executor.py"),
+    ).toBe(true);
+    expect(errors.slice(1)).toMatchInlineSnapshot(`
       [
         {
           "cellId": "Hbol",
+          "kind": "cell",
           "lineNumber": 4,
         },
         {
           "cellId": "Hbol",
+          "kind": "cell",
           "lineNumber": 2,
         },
       ]
@@ -22,10 +28,15 @@ describe("traceback", () => {
 
   test("extracts cell-link from assertion", () => {
     const info = extractAllTracebackInfo(Tracebacks.assertion);
-    expect(info).toMatchInlineSnapshot(`
+    expect(
+      info[0].kind === "file" &&
+        info[0].filePath.endsWith("marimo/_runtime/executor.py"),
+    ).toBe(true);
+    expect(info.slice(1)).toMatchInlineSnapshot(`
       [
         {
           "cellId": "Hbol",
+          "kind": "cell",
           "lineNumber": 1,
         },
       ]

frontend/src/utils/traceback.ts

@@ -28,13 +28,38 @@ export const elementContainsMarimoCellFile = (domNode: Element) => {
   );
 };
 
-export interface TracebackInfo {
-  cellId: CellId;
-  lineNumber: number;
-}
+export type TracebackInfo =
+  | {
+      kind: "file";
+      filePath: string;
+      lineNumber: number;
+    }
+  | {
+      kind: "cell";
+      cellId: CellId;
+      lineNumber: number;
+    };
 
 /**
  * Extract the cell id and line number from a traceback DOM node.
+ *
+ * Example transformation:
+ *
+ *   File <span class="nb">"/tmp/marimo_<number>/__marimo__cell_<CellId>.py"</span>
+ *   , line <span class="n">1</span>...
+ *
+ *   becomes
+ *
+ *   { kind: "cell", cellId: <CellID>, lineNumber: 1 }
+ *
+ *   or for files:
+ *
+ *   File <span class="nb">"/path/to/file.py"</span>
+ *   , line <span class="n">42</span>...
+ *
+ *   becomes
+ *
+ *   { kind: "file", filePath: "/path/to/file.py", lineNumber: 42 }
  */
 export function getTracebackInfo(domNode: DOMNode): TracebackInfo | null {
   // The traceback can be manipulated either in output render or in the pygments
@@ -57,7 +82,7 @@ export function getTracebackInfo(domNode: DOMNode): TracebackInfo | null {
   if (
     domNode instanceof Element &&
     domNode.firstChild instanceof Text &&
-    elementContainsMarimoCellFile(domNode)
+    matchesSelector(domNode, "span.nb")
   ) {
     const nextSibling = domNode.next;
     if (nextSibling && nextSibling instanceof Text) {
@@ -68,15 +93,22 @@ export function getTracebackInfo(domNode: DOMNode): TracebackInfo | null {
         lineSibling.firstChild instanceof Text &&
         matchesSelector(lineSibling, "span.m")
       ) {
-        const cellId = /__marimo__cell_(\w+)_/.exec(
-          domNode.firstChild.nodeValue,
-        )?.[1];
         const lineNumber = Number.parseInt(
           lineSibling.firstChild.nodeValue || "0",
           10,
         );
-        if (cellId && lineNumber) {
-          return { cellId: cellId as CellId, lineNumber };
+        if (domNode.firstChild.nodeValue?.includes("__marimo__")) {
+          const cellId = /__marimo__cell_(\w+)_/.exec(
+            domNode.firstChild.nodeValue,
+          )?.[1] as CellId;
+          if (cellId && lineNumber) {
+            return { kind: "cell", cellId, lineNumber };
+          }
+        } else {
+          const filePath = /"(.+?)"/.exec(domNode.firstChild.nodeValue)?.[1];
+          if (filePath && lineNumber) {
+            return { kind: "file", filePath, lineNumber };
+          }
         }
       }
     }
@@ -93,11 +125,8 @@ export function extractAllTracebackInfo(traceback: string): TracebackInfo[] {
     replace: (domNode) => {
       const info = getTracebackInfo(domNode);
       if (info) {
-        infos.push({
-          cellId: info.cellId,
-          lineNumber: info.lineNumber,
-        });
-        return `${info.cellId}:${info.lineNumber}`;
+        infos.push(info);
+        return "dummy";
       }
     },
   });

test.py

@@ -0,0 +1,26 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "marimo",
+# ]
+# ///
+import marimo
+
+__generated_with = "0.16.2"
+app = marimo.App(width="full")
+
+
+@app.cell
+def _():
+    import pathlib
+    pathlib.Path(5)
+    return
+
+
+@app.cell
+def _():
+    return
+
+
+if __name__ == "__main__":
+    app.run()