Add TaskGraph.keys

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")
@@ -113,6 +113,19 @@ def compute(
         else:
             return self._scheduler.get(self._graph, [keys])[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
     ) -> graphviz.Digraph:  # type: ignore[name-defined] # noqa: F821

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
@@ -53,3 +71,10 @@ def test_compute_raises_when_provided_with_key_not_in_graph() -> None:
         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]}