Add TaskGraph.keys

src/sciline/pipeline.py

@@ -860,7 +860,7 @@ def get(
         else:
             graph = self.build(keys, handler=handler)  # type: ignore[arg-type]
         return TaskGraph(
-            graph=graph, keys=keys, scheduler=scheduler  # type: ignore[arg-type]
+            graph=graph, targets=keys, scheduler=scheduler  # type: ignore[arg-type]
         )
 
     @overload

src/sciline/task_graph.py

@@ -3,10 +3,10 @@
 from __future__ import annotations
 
 from html import escape
-from typing import Any, Optional, Sequence, Tuple, TypeVar, Union
+from typing import Any, Generator, Optional, Sequence, Tuple, TypeVar, Union
 
 from .scheduler import DaskScheduler, NaiveScheduler, Scheduler
-from .typing import Graph, Item
+from .typing import Graph, Item, Key
 from .utils import keyname
 
 T = TypeVar("T")
@@ -70,11 +70,11 @@ def __init__(
         self,
         *,
         graph: Graph,
-        keys: Union[type, Tuple[type, ...], Item[T], Tuple[Item[T], ...]],
+        targets: Union[type, Tuple[type, ...], Item[T], Tuple[Item[T], ...]],
         scheduler: Optional[Scheduler] = None,
     ) -> None:
         self._graph = graph
-        self._keys = keys
+        self._keys = targets
         if scheduler is None:
             try:
                 scheduler = DaskScheduler()
@@ -84,7 +84,7 @@ def __init__(
 
     def compute(
         self,
-        keys: Optional[
+        targets: Optional[
             Union[type, Tuple[type, ...], Item[T], Tuple[Item[T], ...]]
         ] = None,
     ) -> Any:
@@ -93,25 +93,38 @@ def compute(
 
         Parameters
         ----------
-        keys:
+        targets:
             Optional list of keys to compute. This can be used to override the keys
             stored in the graph instance. Note that the keys must be present in the
             graph as intermediate results, otherwise KeyError is raised.
 
         Returns
         -------
-        If ``keys`` is a single type, returns the single result that was computed.
-        If ``keys`` is a tuple of types, returns a dictionary with type as keys
+        If ``targets`` is a single type, returns the single result that was computed.
+        If ``targets`` is a tuple of types, returns a dictionary with type as keys
         and the corresponding results as values.
-
         """
-        if keys is None:
-            keys = self._keys
-        if isinstance(keys, tuple):
-            results = self._scheduler.get(self._graph, list(keys))
-            return dict(zip(keys, results))
+        if targets is None:
+            targets = self._keys
+        if isinstance(targets, tuple):
+            results = self._scheduler.get(self._graph, list(targets))
+            return dict(zip(targets, results))
         else:
-            return self._scheduler.get(self._graph, [keys])[0]
+            return self._scheduler.get(self._graph, [targets])[0]
+
+    def keys(self) -> Generator[Key, None, None]:
+        """
+        Iterate over all keys of the graph.
+
+        Yields all keys, i.e., the types of values that can be computed or are
+        provided as parameters.
+
+        Returns
+        -------
+        :
+            Iterable over keys.
+        """
+        yield from self._graph.keys()
 
     def visualize(
         self, **kwargs: Any

tests/task_graph_test.py

@@ -1,11 +1,29 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
+from typing import NewType, TypeVar
+
 import pytest
 
 import sciline as sl
 from sciline.task_graph import TaskGraph
 from sciline.typing import Graph
 
+A = NewType('A', int)
+B = NewType('B', int)
+T = TypeVar('T', A, B)
+
+
+class Str(sl.Scope[T, str], str):
+    ...
+
+
+def to_string(x: T) -> Str[T]:
+    return Str[T](str(x))
+
+
+def repeat(a: A, s: Str[B]) -> list[str]:
+    return [s] * a
+
 
 def as_float(x: int) -> float:
     return 0.5 * x
@@ -18,38 +36,48 @@ def make_task_graph() -> Graph:
 
 def test_default_scheduler_is_dask_when_dask_available() -> None:
     _ = pytest.importorskip("dask")
-    tg = TaskGraph(graph={}, keys=())
+    tg = TaskGraph(graph={}, targets=())
     assert isinstance(tg._scheduler, sl.scheduler.DaskScheduler)
 
 
 def test_compute_returns_value_when_initialized_with_single_key() -> None:
     graph = make_task_graph()
-    tg = TaskGraph(graph=graph, keys=float)
+    tg = TaskGraph(graph=graph, targets=float)
     assert tg.compute() == 0.5
 
 
 def test_compute_returns_dict_when_initialized_with_key_tuple() -> None:
     graph = make_task_graph()
-    assert TaskGraph(graph=graph, keys=(float,)).compute() == {float: 0.5}
-    assert TaskGraph(graph=graph, keys=(float, int)).compute() == {float: 0.5, int: 1}
+    assert TaskGraph(graph=graph, targets=(float,)).compute() == {float: 0.5}
+    assert TaskGraph(graph=graph, targets=(float, int)).compute() == {
+        float: 0.5,
+        int: 1,
+    }
 
 
 def test_compute_returns_value_when_provided_with_single_key() -> None:
     graph = make_task_graph()
-    tg = TaskGraph(graph=graph, keys=float)
+    tg = TaskGraph(graph=graph, targets=float)
     assert tg.compute(int) == 1
 
 
 def test_compute_returns_dict_when_provided_with_key_tuple() -> None:
     graph = make_task_graph()
-    tg = TaskGraph(graph=graph, keys=float)
+    tg = TaskGraph(graph=graph, targets=float)
     assert tg.compute((int, float)) == {int: 1, float: 0.5}
 
 
 def test_compute_raises_when_provided_with_key_not_in_graph() -> None:
     graph = make_task_graph()
-    tg = TaskGraph(graph=graph, keys=float)
+    tg = TaskGraph(graph=graph, targets=float)
     with pytest.raises(KeyError):
         tg.compute(str)
     with pytest.raises(KeyError):
         tg.compute((str, float))
+
+
+def test_keys_iter() -> None:
+    pl = sl.Pipeline([to_string, repeat], params={A: 3, B: 4})
+    tg = pl.get(list[str])
+    assert len(list(tg.keys())) == 4  # there are no duplicates
+    assert set(tg.keys()) == {A, B, Str[B], list[str]}