21 from __future__
import annotations
22 from collections
import defaultdict
24 __all__ = (
"_DatasetTracker",
"DatasetTypeName")
26 from dataclasses
import dataclass, field
28 from typing
import (DefaultDict, Generic, Optional, Set, TypeVar, Generator, Tuple, NewType,)
30 from lsst.daf.butler
import DatasetRef
32 from .quantumNode
import QuantumNode
33 from ..pipeline
import TaskDef
36 DatasetTypeName = NewType(
"DatasetTypeName", str)
39 _T = TypeVar(
"_T", DatasetTypeName, DatasetRef)
40 _U = TypeVar(
"_U", TaskDef, QuantumNode)
45 inputs: Set[_U] = field(default_factory=set)
46 output: Optional[_U] =
None
51 self._container: DefaultDict[_T, _DatasetTrackerElement[_U]] = defaultdict(_DatasetTrackerElement)
54 self._container[key].inputs.add(value)
57 element = self._container[key]
58 if element.output
is not None:
59 raise ValueError(f
"Only one output for key {key} is allowed, "
60 f
"the current output is set to {element.output}")
61 element.output = value
64 return self._container[key].inputs
67 return self._container[key].output
69 def getAll(self, key: _T) -> Set[_U]:
70 output = self._container[key].output
71 if output
is not None:
72 return self._container[key].inputs.union((output,))
73 return set(self._container[key].inputs)
78 if None in graph.nodes():
79 graph.remove_node(
None)
82 def _datasetDictToEdgeIterator(self) -> Generator[Tuple[Optional[_U], Optional[_U]], None, None]:
83 """Helper function designed to be used in conjunction with
84 `networkx.DiGraph.add_edges_from`. This takes a mapping of keys to
85 `_DatasetTrackers` and yields successive pairs of elements that are to
86 be considered connected by the graph.
88 for entry
in self._container.values():
92 inputs = entry.inputs
or (
None,)
94 yield (entry.output, inpt)
97 return set(self._container.
keys())
Set[_U] getAll(self, _T key)
Set[_U] getInputs(self, _T key)
def addOutput(self, _T key, _U value)
Optional[_U] getOutput(self, _T key)
def addInput(self, _T key, _U value)
Generator[Tuple[Optional[_U], Optional[_U]], None, None] _datasetDictToEdgeIterator(self)
nx.DiGraph makeNetworkXGraph(self)