Repository: marimo-team/marimo
Pull request: Improvement: annotate cycles with variable names
Fixes #847

Problem statement:
Naming in cycle resolution
### Description

I am trying to migrate a notebook with a cycle. It is hard to debug especially when the cycle is long.

### Suggested solution

Could the edges of the dependency graph be annotated with the variable names? For example 
```
This cell is in a cycle:
cell-5 -> cell-0
cell-0 -> cell-5
```

Change to 
`cell-5 (var a...) -> cell-0`?
Or it could be shown graphically in the dependency explorer, on the edges, maybe even just on hover?

### Alternative

_No response_

### Additional context

_No response_

Hints from issue discussion:
Fantastic suggestion -- we should do this! Will get around to it when we have spare time.

Commit message: switched to list instead of set

Patch to review:
diff --git a/frontend/src/components/editor/output/MarimoErrorOutput.tsx b/frontend/src/components/editor/output/MarimoErrorOutput.tsx
index 80ea506f963..73928b4c21b 100644
--- a/frontend/src/components/editor/output/MarimoErrorOutput.tsx
+++ b/frontend/src/components/editor/output/MarimoErrorOutput.tsx
@@ -63,11 +63,12 @@ export const MarimoErrorOutput = ({
           <Fragment key={idx}>
             <p className="mt-4">{"This cell is in a cycle:"}</p>
             <ul className="list-disc">
-              {error.edges.map((edge) => (
+              {error.edges_with_vars.map((edge) => (
                 <li className={liStyle} key={`${edge[0]}-${edge[1]}`}>
                   <CellLinkError cellId={edge[0]} />
+                  {" -> "} {edge[1].length == 1 ? edge[1] : edge[1].join(", ")}{" "}
                   {" -> "}
-                  <CellLinkError cellId={edge[1]} />
+                  <CellLinkError cellId={edge[2]} />
                 </li>
               ))}
             </ul>
diff --git a/frontend/src/core/kernel/messages.ts b/frontend/src/core/kernel/messages.ts
index 4af5cb9ae18..3ad69b5963c 100644
--- a/frontend/src/core/kernel/messages.ts
+++ b/frontend/src/core/kernel/messages.ts
@@ -26,7 +26,7 @@ export type MarimoError =
       raising_cell?: CellId;
     }
   | { type: "ancestor-stopped"; msg: string; raising_cell: CellId }
-  | { type: "cycle"; edges: Array<[CellId, CellId]> }
+  | { type: "cycle"; edges_with_vars: Array<[CellId, string[], CellId]> }
   | { type: "multiple-defs"; name: string; cells: CellId[] }
   | { type: "delete-nonlocal"; name: string; cells: CellId[] }
   | { type: "unknown"; msg?: string };
diff --git a/marimo/_messaging/errors.py b/marimo/_messaging/errors.py
index 05ebf7759c8..fd5bb3200f9 100644
--- a/marimo/_messaging/errors.py
+++ b/marimo/_messaging/errors.py
@@ -5,12 +5,12 @@
 from typing import Optional, Union
 
 from marimo._ast.cell import CellId_t
-from marimo._runtime.dataflow import Edge
+from marimo._runtime.dataflow import EdgeWithVar
 
 
 @dataclass
 class CycleError:
-    edges: tuple[Edge, ...]
+    edges_with_vars: tuple[EdgeWithVar, ...]
     type: str = "cycle"
 
 
diff --git a/marimo/_runtime/dataflow.py b/marimo/_runtime/dataflow.py
index 15441a07c75..46c6c8830c4 100644
--- a/marimo/_runtime/dataflow.py
+++ b/marimo/_runtime/dataflow.py
@@ -3,7 +3,7 @@
 
 import threading
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Optional, Tuple
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple
 
 from marimo import _loggers
 from marimo._ast.cell import (
@@ -19,6 +19,9 @@
     from collections.abc import Collection
 
 Edge = Tuple[CellId_t, CellId_t]
+# EdgeWithVar uses a list rather than a set for the variables linking the cells
+#   as sets are not JSON-serializable (required by static_notebook_template())
+EdgeWithVar = Tuple[CellId_t, List[str], CellId_t]
 
 LOGGER = _loggers.marimo_logger()
 
diff --git a/marimo/_runtime/validate_graph.py b/marimo/_runtime/validate_graph.py
index da0ba599e36..82a6b59dd63 100644
--- a/marimo/_runtime/validate_graph.py
+++ b/marimo/_runtime/validate_graph.py
@@ -57,8 +57,19 @@ def check_for_cycles(graph: DirectedGraph) -> dict[CellId_t, list[CycleError]]:
     errors = defaultdict(list)
     for cycle in graph.cycles:
         nodes_in_cycle: set[CellId_t] = set(sum(cycle, ()))
+        # before reporting the cells in the cycle to the user,
+        #   we first annotate the cycle with the variable names
+        #   that link its cells
+        cycle_with_vars = tuple(
+            (
+                edge[0],
+                list(graph.cells[edge[0]].defs & graph.cells[edge[1]].refs),
+                edge[1],
+            )
+            for edge in cycle
+        )
         for cid in nodes_in_cycle:
-            errors[cid].append(CycleError(edges=cycle))
+            errors[cid].append(CycleError(edges_with_vars=cycle_with_vars))
     return errors
 
 
diff --git a/tests/_runtime/test_runtime.py b/tests/_runtime/test_runtime.py
index c5ec3a62979..888a039d321 100644
--- a/tests/_runtime/test_runtime.py
+++ b/tests/_runtime/test_runtime.py
@@ -15,7 +15,7 @@
 )
 from marimo._messaging.types import NoopStream
 from marimo._plugins.ui._core.ids import IDProvider
-from marimo._runtime.dataflow import Edge
+from marimo._runtime.dataflow import EdgeWithVar
 from marimo._runtime.requests import (
     AppMetadata,
     CreationRequest,
@@ -32,11 +32,11 @@
     from types import ModuleType
 
 
-def _check_edges(error: Error, expected_edges: Sequence[Edge]) -> None:
+def _check_edges(error: Error, expected_edges: Sequence[EdgeWithVar]) -> None:
     assert isinstance(error, CycleError)
-    assert len(error.edges) == len(expected_edges)
+    assert len(error.edges_with_vars) == len(expected_edges)
     for edge in expected_edges:
-        assert edge in error.edges or (edge[1], edge[0]) in error.edges
+        assert edge in error.edges_with_vars or reversed(edge) in error.edges_with_vars
 
 
 # Test Basic Reactivity
@@ -335,8 +335,8 @@ async def test_cycle_error(k: Kernel) -> None:
     assert set(k.errors.keys()) == {"0", "1"}
     assert len(k.errors["0"]) == 1
     assert len(k.errors["1"]) == 1
-    _check_edges(k.errors["0"][0], [("0", "1"), ("1", "0")])
-    _check_edges(k.errors["1"][0], [("0", "1"), ("1", "0")])
+    _check_edges(k.errors["0"][0], [("0", ["x"], "1"), ("1", ["y"], "0")])
+    _check_edges(k.errors["1"][0], [("0", ["x"], "1"), ("1", ["y"], "0")])
 
     # break cycle by modifying cell
     await k.run([ExecutionRequest(cell_id="1", code="y=1")])
@@ -358,7 +358,7 @@ async def test_break_cycle_error_with_delete(k: Kernel) -> None:
     assert set(k.errors.keys()) == {"0", "1"}
     assert len(k.errors["0"]) == 1
     assert len(k.errors["1"]) == 1
-    _check_edges(k.errors["0"][0], [("0", "1"), ("1", "0")])
+    _check_edges(k.errors["0"][0], [("0", ["x"], "1"), ("1", ["y"], "0")])
 
     # break cycle by deleting cell
     await k.delete(DeleteRequest(cell_id="1"))
diff --git a/tests/_runtime/test_validate_graph.py b/tests/_runtime/test_validate_graph.py
index e73a10bd2e2..730b8f6f59b 100644
--- a/tests/_runtime/test_validate_graph.py
+++ b/tests/_runtime/test_validate_graph.py
@@ -83,8 +83,8 @@ def test_two_node_cycle() -> None:
     graph.register_cell("1", parse_cell("y = x"))
     errors = check_for_errors(graph)
     assert set(errors.keys()) == set(["0", "1"])
-    assert errors["0"] == (CycleError(edges=(("0", "1"), ("1", "0"))),)
-    assert errors["1"] == (CycleError(edges=(("0", "1"), ("1", "0"))),)
+    assert errors["0"] == (CycleError(edges_with_vars=(("0", ["x"], "1"), ("1", ["y"], "0"))),)
+    assert errors["1"] == (CycleError(edges_with_vars=(("0", ["x"], "1"), ("1", ["y"], "0"))),)
 
 
 def test_three_node_cycle() -> None:
@@ -97,11 +97,11 @@ def test_three_node_cycle() -> None:
     for _, t in errors.items():
         assert len(t) == 1
         assert isinstance(t[0], CycleError)
-        edges = t[0].edges
-        assert len(edges) == 3
-        assert ("0", "1") in edges or ("1", "0") in edges
-        assert ("1", "2") in edges or ("2", "1") in edges
-        assert ("0", "2") in edges or ("2", "0") in edges
+        edges_with_vars = t[0].edges_with_vars
+        assert len(edges_with_vars) == 3
+        assert ("0", ["x"], "1") in edges_with_vars or ("1", ["y"], "0") in edges_with_vars
+        assert ("1", ["y"], "2") in edges_with_vars or ("2", ["z"], "1") in edges_with_vars
+        assert ("2", ["z"], "0") in edges_with_vars or ("0", ["x"], "2") in edges_with_vars
 
 
 def test_cycle_and_multiple_def() -> None:
@@ -119,10 +119,10 @@ def test_cycle_and_multiple_def() -> None:
         cycle_error = cast(
             CycleError, t[0] if isinstance(t[0], CycleError) else t[1]
         )
-        edges = cycle_error.edges
-        assert len(edges) == 2
-        assert ("0", "1") in edges
-        assert ("1", "0") in edges
+        edges_with_vars = cycle_error.edges_with_vars
+        assert len(edges_with_vars) == 2
+        assert ("0", ["x"], "1") in edges_with_vars
+        assert ("1", ["y"], "0") in edges_with_vars
 
         multiple_definition_error = cast(
             MultipleDefinitionError,


Code review comments:
diff_hunk: @@ -32,11 +32,11 @@
     from types import ModuleType
 
 
-def _check_edges(error: Error, expected_edges: Sequence[Edge]) -> None:
+def _check_edges(error: Error, expected_edges: Sequence[EdgeWithVar]) -> None:
     assert isinstance(error, CycleError)
-    assert len(error.edges) == len(expected_edges)
+    assert len(error.edges_with_vars) == len(expected_edges)
     for edge in expected_edges:
-        assert edge in error.edges or (edge[1], edge[0]) in error.edges
+        assert edge in error.edges_with_vars or reversed(edge) in error.edges_with_vars
path: tests/_runtime/test_runtime.py
text: @user1:
I believe the reversed logic wasn't needed originally; removed it and these tests
still pass on my machine. Checking for the presence of a reversed edge is no longer
correct, since the variables that flow on the edges change when you flip the direction.

```suggestion
        assert edge in error.edges_with_vars
```

diff_hunk: @@ -63,11 +63,12 @@ export const MarimoErrorOutput = ({
           <Fragment key={idx}>
             <p className="mt-4">{"This cell is in a cycle:"}</p>
             <ul className="list-disc">
-              {error.edges.map((edge) => (
+              {error.edges_with_vars.map((edge) => (
                 <li className={liStyle} key={`${edge[0]}-${edge[1]}`}>
                   <CellLinkError cellId={edge[0]} />
+                  {" -> "} {edge[1].length == 1 ? edge[1] : edge[1].join(", ")}{" "}
path: frontend/src/components/editor/output/MarimoErrorOutput.tsx
text: @user1:
```suggestion
                   <span className="text-muted-foreground">                       
                      {" - "} {edge[1].length == 1 ? edge[1] : edge[1].join(", ")} 
                      {" -> "}                                                     
                   </span>  
```