feat: basic Debugpy integration

marimo/_ast/app.py

@@ -4,6 +4,7 @@
 import ast
 import base64
 import inspect
+import os
 import sys
 import threading
 from collections.abc import Iterable, Iterator, Mapping
@@ -594,6 +595,19 @@ def _globals_to_defs(self, glbls: dict[str, Any]) -> _Namespace:
     def run(
         self,
     ) -> tuple[Sequence[Any], Mapping[str, Any]]:
+        # Enabled specifically for debugging purposes.
+        # see docs.marimo.io/guides/debugging
+        if os.environ.get("MARIMO_SCRIPT_EDIT"):
+            from marimo._cli.cli import edit
+
+            if self._filename is None:
+                raise RuntimeError(
+                    "MARIMO_SCRIPT_EDIT is set, but filename cannot be determined."
+                )
+            ctx = edit.make_context("edit", ["--watch", self._filename])
+            edit.invoke(ctx)
+            return ((), {})
+
         self._maybe_initialize()
         glbls: dict[str, Any] = {}
         if self._setup is not None:

marimo/_ast/cell_manager.py

@@ -227,7 +227,8 @@ def register_ir_cell(
             cell_id = CellId_t(SETUP_CELL_NAME)
         else:
             cell_id = self.create_cell_id()
-        cell = ir_cell_factory(cell_def, cell_id=cell_id)
+        filename = app.filename if app is not None else None
+        cell = ir_cell_factory(cell_def, cell_id=cell_id, filename=filename)
         cell_config = CellConfig.from_dict(
             cell_def.options,
         )

marimo/_ast/compiler.py

@@ -152,7 +152,16 @@ def compile_cell(
     source_position: Optional[SourcePosition] = None,
     carried_imports: list[ImportData] | None = None,
     test_rewrite: bool = False,
+    filename: Optional[str] = None,
 ) -> CellImpl:
+    if filename is not None and source_position is None:
+        source_position = solve_source_position(
+            code,
+            filename,
+        )
+    elif filename is not None and source_position is not None:
+        source_position.filename = filename
+
     # Replace non-breaking spaces with regular spaces -- some frontends
     # send nbsp in place of space, which is a syntax error.
     #
@@ -213,7 +222,6 @@ def compile_cell(
     expr.col_offset = final_expr.end_col_offset  # type: ignore[attr-defined]
     expr.end_col_offset = final_expr.end_col_offset  # type: ignore[attr-defined]
 
-    filename: str
     if source_position:
         # Modify the "source" position for meaningful stacktraces
         fix_source_position(module, source_position)
@@ -294,6 +302,36 @@ def compile_cell(
     )
 
 
+def solve_source_position(
+    code: str, filename: str
+) -> Optional[SourcePosition]:
+    from marimo._ast.load import _maybe_contents
+    from marimo._ast.parse import parse_notebook
+    from marimo._utils.cell_matching import match_cell_ids_by_similarity
+
+    contents = _maybe_contents(filename)
+    if not contents:
+        return None
+
+    notebook = parse_notebook(contents)
+    if notebook is None or not notebook.valid:
+        return None
+    on_disk = {
+        CellId_t(str(i)): cell.code for i, cell in enumerate(notebook.cells)
+    }
+    matches = match_cell_ids_by_similarity(on_disk, {CellId_t("new"): code})
+    if not matches or len(matches) != 1:
+        return None
+    (cell_index,) = matches.keys()
+    index = int(cell_index)
+
+    return SourcePosition(
+        filename=filename,
+        lineno=notebook.cells[index].lineno,
+        col_offset=notebook.cells[index].col_offset,
+    )
+
+
 def get_source_position(
     f: Cls | Callable[..., Any], lineno: int, col_offset: int
 ) -> Optional[SourcePosition]:
@@ -427,9 +465,21 @@ def toplevel_cell_factory(
     )
 
 
-def ir_cell_factory(cell_def: CellDef, cell_id: CellId_t) -> Cell:
+def ir_cell_factory(
+    cell_def: CellDef, cell_id: CellId_t, filename: Optional[str] = None
+) -> Cell:
     # NB. no need for test rewrite, anonymous file, etc.
     # Because this is never invoked in script mode.
+    source_position = None
+    # EXCEPT in the case of debugpy, where we need to preserve source position.
+    if os.environ.get("DEBUGPY_RUNNING"):
+        if filename and cell_def.lineno:
+            source_position = SourcePosition(
+                filename=filename,
+                lineno=cell_def.lineno,
+                col_offset=cell_def.col_offset,
+            )
+
     prefix = ""
     if isinstance(cell_def, (FunctionCell, ClassCell)):
         prefix = TOPLEVEL_CELL_PREFIX
@@ -438,6 +488,7 @@ def ir_cell_factory(cell_def: CellDef, cell_id: CellId_t) -> Cell:
         _cell=compile_cell(
             cell_def.code,
             cell_id=cell_id,
+            source_position=source_position,
         ),
     )
 

marimo/_runtime/dataflow.py

@@ -594,6 +594,32 @@ def get_transitive_references(
             return processed | refs
         return processed - refs
 
+    def copy(self, filename: None | str = None) -> DirectedGraph:
+        """Return a deep copy of the graph by recompiling all cells.
+
+        This is mainly useful in the case where recompilation must be done
+        due to a dynamically changing notebook, where the line cache must be
+        consistent with the cell code, e.g. for debugging.
+        """
+        from marimo._ast.compiler import compile_cell
+
+        graph = DirectedGraph()
+        with self.lock:
+            for cid, old_cell in self.cells.items():
+                cell = compile_cell(
+                    old_cell.code,
+                    cell_id=cid,
+                    filename=filename,
+                )
+                # Carry over import data manually
+                imported_defs = old_cell.import_workspace.imported_defs
+                is_import_block = old_cell.import_workspace.is_import_block
+                cell.import_workspace.imported_defs = imported_defs
+                cell.import_workspace.is_import_block = is_import_block
+                # Reregister
+                graph.register_cell(cid, cell)
+        return graph
+
 
 def transitive_closure(
     graph: DirectedGraph,

marimo/_runtime/runtime.py

@@ -845,7 +845,9 @@ def _try_compiling_cell(
         error: Optional[Error] = None
         try:
             cell = compile_cell(
-                code, cell_id=cell_id, carried_imports=carried_imports
+                code,
+                cell_id=cell_id,
+                carried_imports=carried_imports,
             )
         except Exception as e:
             cell = None
@@ -1395,9 +1397,15 @@ def note_time_of_interruption(
             if isinstance(run_result.exception, MarimoInterrupt):
                 self.last_interrupt_timestamp = time.time()
 
+        # Rebuild graph with sourceful positions
+        # Note, this is relatively expensive, but a reasonable tradeoff
+        graph = self.graph
+        if os.getenv("DEBUGPY_RUNNING"):
+            graph = self.graph.copy(self.app_metadata.filename)
+
         runner = cell_runner.Runner(
             roots=roots,
-            graph=self.graph,
+            graph=graph,
             glbls=self.globals,
             excluded_cells=set(self.errors.keys()),
             debugger=self.debugger,
@@ -2747,7 +2755,7 @@ def launch_kernel(
     stdin = ThreadSafeStdin(stream) if is_edit_mode else None
     debugger = (
         marimo_pdb.MarimoPdb(stdout=stdout, stdin=stdin)
-        if is_edit_mode
+        if is_edit_mode and not bool(os.getenv("DEBUGPY_RUNNING"))
         else None
     )
 

marimo/_smoke_tests/pdb_test.py

@@ -1,20 +1,58 @@
 # Copyright 2024 Marimo. All rights reserved.
+
 import marimo
 
-__generated_with = "0.1.77"
+__generated_with = "0.15.2"
 app = marimo.App()
 
 
 @app.cell
-def __():
+def _():
+    def percentile(xs, p):
+        """
+        Return the p-th percentile of xs.
+        """
+        xs = sorted(xs)
+        idx = round(p / 100 * len(xs))  # here there be bugs
+        return xs[idx]
+
+
+    percentile([1, 2, 3], 100)
+    return
+
+
+@app.cell
+def _():
+    # Compute triangle numbers
+    triangle = 0
+    triangle_count = 20
+    for i in range(1, triangle_count):
+        triangle += i  # T_i = sum of 1..i
+        # Debug at the 10th iteration
+        # as a sanity check. Should be 55.
+        if i == 10:
+            breakpoint()
+    return
+
+
+@app.cell
+def _():
     import marimo as mo
-    return mo,
+    return
 
 
 @app.cell
-def __(mo):
-    # import pdb; pdb.set_trace()
-    mo.pdb.set_trace()
+def _():
+    import pdb
+
+    pdb.set_trace()
+    return
+
+
+@app.cell
+def _():
+
+
     return
 
 

tests/_ast/test_app.py

@@ -1,12 +1,15 @@
 # Copyright 2024 Marimo. All rights reserved.
 
 from __future__ import annotations
+import os
 import pathlib
 import subprocess
 import sys
 import textwrap
 from typing import TYPE_CHECKING, Any
+from unittest.mock import patch
 
+import click
 import pytest
 
 from marimo._ast.app import (