Avoid recursion and allow building cyclic graphs

requirements/test.in

@@ -1,3 +1,4 @@
 -r base.in
+graphviz
 numpy
 pytest

requirements/test.txt

@@ -1,4 +1,4 @@
-# SHA1:3720d9b18830e4fcedb827ec36f6808035c1ea2c
+# SHA1:fa933c4a7cbdd9d4d24d05397dc19e37685ca4a1
 #
 # This file is autogenerated by pip-compile-multi
 # To update, run:
@@ -8,6 +8,8 @@
 -r base.txt
 exceptiongroup==1.1.2
     # via pytest
+graphviz==0.20.1
+    # via -r test.in
 iniconfig==2.0.0
     # via pytest
 numpy==1.24.4

src/sciline/pipeline.py

@@ -2,7 +2,6 @@
 # Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
 from __future__ import annotations
 
-from graphlib import CycleError
 from typing import (
     Any,
     Callable,
@@ -191,20 +190,22 @@ def build(self, tp: Type[T], /) -> Graph:
         tp:
             Type to build the graph for.
         """
-        provider: Callable[..., T]
-        provider, bound = self._get_provider(tp)
-        tps = get_type_hints(provider)
-        args = {
-            name: _bind_free_typevars(t, bound=bound)
-            for name, t in tps.items()
-            if name != 'return'
-        }
-        graph: Graph = {tp: (provider, args)}
-        try:
+        graph: Graph = {}
+        stack: List[Type[T]] = [tp]
+        while stack:
+            tp = stack.pop()
+            provider: Callable[..., T]
+            provider, bound = self._get_provider(tp)
+            tps = get_type_hints(provider)
+            args = {
+                name: _bind_free_typevars(t, bound=bound)
+                for name, t in tps.items()
+                if name != 'return'
+            }
+            graph[tp] = (provider, args)
             for arg in args.values():
-                graph.update(self.build(arg))
-        except RecursionError:
-            raise CycleError("Cycle detected while building graph for", tp)
+                if arg not in graph:
+                    stack.append(arg)
         return graph
 
     @overload

src/sciline/scheduler.py

@@ -9,8 +9,18 @@
 ]
 
 
+class CycleError(Exception):
+    pass
+
+
 class Scheduler(Protocol):
     def get(self, graph: Graph, keys: List[type]) -> Any:
+        """
+        Compute the result for given keys from the graph.
+
+        Must raise :py:class:`sciline.scheduler.CycleError` if the graph contains
+        a cycle.
+        """
         ...
 
 
@@ -24,12 +34,17 @@ class NaiveScheduler:
     """
 
     def get(self, graph: Graph, keys: List[type]) -> Any:
-        from graphlib import TopologicalSorter
+        import graphlib
 
-        results: Dict[type, Any] = {}
         dependencies = {tp: set(args.values()) for tp, (_, args) in graph.items()}
-        ts = TopologicalSorter(dependencies)
-        for t in ts.static_order():
+        ts = graphlib.TopologicalSorter(dependencies)
+        try:
+            # Create list from generator to force early exception if there is a cycle
+            tasks = list(ts.static_order())
+        except graphlib.CycleError as e:
+            raise CycleError from e
+        results: Dict[type, Any] = {}
+        for t in tasks:
             provider, args = graph[t]
             args = {name: results[arg] for name, arg in args.items()}
             results[t] = provider(*args.values())
@@ -57,4 +72,9 @@ def __init__(self, scheduler: Optional[Callable[..., Any]] = None) -> None:
 
     def get(self, graph: Graph, keys: List[type]) -> Any:
         dsk = {tp: (provider, *args.values()) for tp, (provider, args) in graph.items()}
-        return self._dask_get(dsk, keys)
+        try:
+            return self._dask_get(dsk, keys)
+        except RuntimeError as e:
+            if str(e).startswith("Cycle detected"):
+                raise CycleError from e
+            raise

tests/pipeline_test.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
 from dataclasses import dataclass
-from graphlib import CycleError
 from typing import Generic, List, NewType, TypeVar
 
 import pytest
@@ -492,16 +491,27 @@ class A(sl.Scope[T], int):
     assert pl.compute(A[str]) == A(2)
 
 
-def test_building_graph_with_loop_raises_CycleError() -> None:
+def test_building_graph_with_cycle_succeeds() -> None:
     def f(x: int) -> float:
         return float(x)
 
     def g(x: float) -> int:
         return int(x)
 
     pipeline = sl.Pipeline([f, g])
-    with pytest.raises(CycleError):
-        pipeline.build(int)
+    _ = pipeline.build(int)
+
+
+def test_computing_graph_with_cycle_raises_CycleError() -> None:
+    def f(x: int) -> float:
+        return float(x)
+
+    def g(x: float) -> int:
+        return int(x)
+
+    pipeline = sl.Pipeline([f, g])
+    with pytest.raises(sl.scheduler.CycleError):
+        pipeline.compute(int)
 
 
 def test_get_with_single_key_return_task_graph_that_computes_value() -> None:

tests/visualize_test.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
+import sciline as sl
+from sciline.visualize import to_graphviz
+
+
+def test_can_visualize_graph_with_cycle() -> None:
+    def int_to_float(x: int) -> float:
+        return float(x)
+
+    def float_to_int(x: float) -> int:
+        return int(x)
+
+    pipeline = sl.Pipeline([int_to_float, float_to_int])
+    graph = pipeline.build(int)
+    to_graphviz(graph)