Store DAG payloads directly

Author: tobiasraabe

src/_pytask/dag.py

@@ -17,7 +17,6 @@
 from _pytask.console import format_task_name
 from _pytask.console import render_to_string
 from _pytask.dag_graph import DAG
-from _pytask.dag_graph import DAGNode
 from _pytask.dag_graph import NoCycleError
 from _pytask.dag_graph import find_cycle
 from _pytask.exceptions import ResolvingDependenciesError
@@ -67,7 +66,7 @@ def _create_dag_from_tasks(tasks: list[PTask]) -> DAG:
     """Create the DAG from tasks, dependencies and products."""
 
     def _add_node_data(dag: DAG, node: PNode | PProvisionalNode) -> None:
-        dag.add_node(node.signature, DAGNode.from_node(node))
+        dag.add_node(node.signature, node)
         if isinstance(node, PythonNode) and isinstance(node.value, PythonNode):
             _add_node_data(dag, node.value)
 
@@ -88,7 +87,7 @@ def _add_product(dag: DAG, task: PTask, node: PNode | PProvisionalNode) -> None:
     dag = DAG()
 
     for task in tasks:
-        dag.add_node(task.signature, DAGNode.from_task(task))
+        dag.add_node(task.signature, task)
         tree_map(lambda x: _add_node_data(dag, x), task.depends_on)
         tree_map(lambda x: _add_node_data(dag, x), task.produces)
 
@@ -147,7 +146,7 @@ def _format_cycles(dag: DAG, cycles: list[tuple[str, str]]) -> str:
 
     lines: list[str] = []
     for x in chain:
-        node = dag.nodes[x].value
+        node = dag.nodes[x]
         if isinstance(node, PTask):
             short_name = format_task_name(node, editor_url_scheme="no_link").plain
         elif isinstance(node, (PNode, PProvisionalNode)):
@@ -173,17 +172,19 @@ def _check_if_tasks_have_the_same_products(dag: DAG, paths: list[Path]) -> None:
     nodes_created_by_multiple_tasks = []
 
     for node in dag.nodes:
-        if dag.nodes[node].node is not None:
+        if isinstance(dag.nodes[node], (PNode, PProvisionalNode)):
             parents = list(dag.predecessors(node))
             if len(parents) > 1:
                 nodes_created_by_multiple_tasks.append(node)
 
     if nodes_created_by_multiple_tasks:
         dictionary = {}
         for node in nodes_created_by_multiple_tasks:
-            short_node_name = format_node_name(
-                dag.nodes[node].node_or_raise(), paths
-            ).plain
+            payload = dag.nodes[node]
+            if not isinstance(payload, (PNode, PProvisionalNode)):
+                msg = f"Expected product node for signature {node!r}."
+                raise TypeError(msg)
+            short_node_name = format_node_name(payload, paths).plain
             short_predecessors = reduce_names_of_multiple_nodes(
                 dag.predecessors(node), dag, paths
             )

src/_pytask/dag_graph.py

@@ -10,79 +10,37 @@
 from typing import cast
 
 from _pytask.compat import import_optional_dependency
+from _pytask.node_protocols import PNode
+from _pytask.node_protocols import PProvisionalNode
+from _pytask.node_protocols import PTask
 
 if TYPE_CHECKING:
     from collections.abc import Callable
     from collections.abc import Iterable
     from collections.abc import Iterator
     from collections.abc import Mapping
 
-    from _pytask.node_protocols import PNode
-    from _pytask.node_protocols import PProvisionalNode
-    from _pytask.node_protocols import PTask
+
+DAGEntry = PTask | PNode | PProvisionalNode
 
 
 class NoCycleError(Exception):
     """Raised when no cycle is found in a graph."""
 
 
-@dataclass(slots=True)
-class DAGNode:
-    """Payload stored for nodes in pytask's internal DAG."""
-
-    task: PTask | None = None
-    node: PNode | PProvisionalNode | None = None
-
-    def __post_init__(self) -> None:
-        if (self.task is None) == (self.node is None):
-            msg = "A DAG node must store exactly one of 'task' or 'node'."
-            raise ValueError(msg)
-
-    @classmethod
-    def from_task(cls, task: PTask) -> DAGNode:
-        """Create a DAG node from a task."""
-        return cls(task=task)
-
-    @classmethod
-    def from_node(cls, node: PNode | PProvisionalNode) -> DAGNode:
-        """Create a DAG node from a dependency or product node."""
-        return cls(node=node)
-
-    @property
-    def value(self) -> PTask | PNode | PProvisionalNode:
-        """Return the wrapped task or node."""
-        if self.task is not None:
-            return self.task
-        return cast("PNode | PProvisionalNode", self.node)
-
-    def task_or_raise(self) -> PTask:
-        """Return the wrapped task."""
-        if self.task is None:
-            msg = "Expected DAG payload to contain a task."
-            raise TypeError(msg)
-        return self.task
-
-    def node_or_raise(self) -> PNode | PProvisionalNode:
-        """Return the wrapped dependency or product node."""
-        if self.node is None:
-            msg = "Expected DAG payload to contain a node."
-            raise TypeError(msg)
-        return self.node
-
-
 @dataclass
 class DAG:
     """A minimal directed graph tailored to pytask's needs."""
 
-    _node_data: dict[str, DAGNode] = field(default_factory=dict)
+    _node_data: dict[str, DAGEntry] = field(default_factory=dict)
     _successors: dict[str, set[str]] = field(default_factory=dict)
     _predecessors: dict[str, set[str]] = field(default_factory=dict)
 
     @property
-    def nodes(self) -> dict[str, DAGNode]:
+    def nodes(self) -> dict[str, DAGEntry]:
         return self._node_data
 
-    def add_node(self, node_name: str, data: DAGNode) -> None:
+    def add_node(self, node_name: str, data: DAGEntry) -> None:
         if node_name not in self._node_data:
             self._successors[node_name] = set()
             self._predecessors[node_name] = set()

src/_pytask/dag_utils.py

@@ -13,18 +13,17 @@
 from _pytask.dag_graph import descendants
 from _pytask.dag_graph import find_cycle
 from _pytask.mark_utils import has_mark
+from _pytask.node_protocols import PTask
 
 if TYPE_CHECKING:
     from collections.abc import Generator
     from collections.abc import Iterable
 
-    from _pytask.node_protocols import PTask
-
 
 def descending_tasks(task_name: str, dag: DAG) -> Generator[str, None, None]:
     """Yield only descending tasks."""
     for descendant in descendants(dag, task_name):
-        if dag.nodes[descendant].task is not None:
+        if isinstance(dag.nodes[descendant], PTask):
             yield descendant
 
 
@@ -37,7 +36,7 @@ def task_and_descending_tasks(task_name: str, dag: DAG) -> Generator[str, None,
 def preceding_tasks(task_name: str, dag: DAG) -> Generator[str, None, None]:
     """Yield only preceding tasks."""
     for ancestor in ancestors(dag, task_name):
-        if dag.nodes[ancestor].task is not None:
+        if isinstance(dag.nodes[ancestor], PTask):
             yield ancestor
 
 
@@ -86,7 +85,7 @@ def from_dag(cls, dag: DAG) -> TopologicalSorter:
         """Instantiate from a DAG."""
         cls.check_dag(dag)
 
-        tasks = [node.task for node in dag.nodes.values() if node.task is not None]
+        tasks = [node for node in dag.nodes.values() if isinstance(node, PTask)]
         priorities = _extract_priorities_from_tasks(tasks)
 
         task_signatures = {task.signature for task in tasks}

src/_pytask/database_utils.py

@@ -116,7 +116,7 @@ def update_states_in_database(session: Session, task_signature: str) -> None:
     if _ENGINE is None:
         return
     for name in node_and_neighbors(session.dag, task_signature):
-        node = session.dag.nodes[name].value
+        node = session.dag.nodes[name]
         if isinstance(node, PProvisionalNode):
             continue
         hash_ = node.state()

src/_pytask/execute.py

@@ -89,7 +89,10 @@ def pytask_execute_build(session: Session) -> bool | None:
     if isinstance(session.scheduler, TopologicalSorter):
         while session.scheduler.is_active():
             task_name = session.scheduler.get_ready()[0]
-            task = session.dag.nodes[task_name].task_or_raise()
+            task = session.dag.nodes[task_name]
+            if not isinstance(task, PTask):
+                msg = f"Expected task node for signature {task_name!r}."
+                raise TypeError(msg)
             report = session.hook.pytask_execute_task_protocol(
                 session=session, task=task
             )
@@ -172,7 +175,7 @@ def pytask_execute_task_setup(session: Session, task: PTask) -> None:  # noqa: C
     if not needs_to_be_executed:
         predecessors = set(dag.predecessors(task.signature)) | {task.signature}
         for node_signature in node_and_neighbors(dag, task.signature):
-            node = dag.nodes[node_signature].value
+            node = dag.nodes[node_signature]
 
             # Skip provisional nodes that are products since they do not have a state.
             if node_signature not in predecessors and isinstance(
@@ -253,7 +256,10 @@ def pytask_execute_task_setup(session: Session, task: PTask) -> None:  # noqa: C
     # Create directory for product if it does not exist. Maybe this should be a `setup`
     # method for the node classes.
     for product in dag.successors(task.signature):
-        node = dag.nodes[product].node_or_raise()
+        node = dag.nodes[product]
+        if not isinstance(node, (PNode, PProvisionalNode)):
+            msg = f"Expected product node for signature {product!r}."
+            raise TypeError(msg)
         if isinstance(node, PPathNode):
             node.path.parent.mkdir(parents=True, exist_ok=True)
         if isinstance(node, DirectoryNode) and node.root_dir:
@@ -347,7 +353,12 @@ def pytask_execute_task_process_report(
         report.outcome = TaskOutcome.WOULD_BE_EXECUTED
 
         for descending_task_name in descending_tasks(task.signature, session.dag):
-            descending_task = session.dag.nodes[descending_task_name].task_or_raise()
+            descending_task = session.dag.nodes[descending_task_name]
+            if not isinstance(descending_task, PTask):
+                msg = (
+                    f"Expected descending task for signature {descending_task_name!r}."
+                )
+                raise TypeError(msg)
             descending_task.markers.append(
                 Mark(
                     "would_be_executed",
@@ -360,7 +371,12 @@ def pytask_execute_task_process_report(
             )
     else:
         for descending_task_name in descending_tasks(task.signature, session.dag):
-            descending_task = session.dag.nodes[descending_task_name].task_or_raise()
+            descending_task = session.dag.nodes[descending_task_name]
+            if not isinstance(descending_task, PTask):
+                msg = (
+                    f"Expected descending task for signature {descending_task_name!r}."
+                )
+                raise TypeError(msg)
             descending_task.markers.append(
                 Mark(
                     "skip_ancestor_failed",

src/_pytask/lockfile.py

@@ -225,7 +225,7 @@ def _build_task_entry(session: Session, task: PTask, root: Path) -> _TaskEntry |
 
     depends_on: dict[str, str] = {}
     for node_signature in predecessors:
-        node = dag.nodes[node_signature].value
+        node = dag.nodes[node_signature]
         if not isinstance(node, (PNode, PTask)):
             continue
         state = node.state()
@@ -240,7 +240,7 @@ def _build_task_entry(session: Session, task: PTask, root: Path) -> _TaskEntry |
 
     produces: dict[str, str] = {}
     for node_signature in successors:
-        node = dag.nodes[node_signature].value
+        node = dag.nodes[node_signature]
         if not isinstance(node, (PNode, PTask)):
             continue
         state = node.state()

src/_pytask/persist.py

@@ -51,7 +51,7 @@ def pytask_execute_task_setup(session: Session, task: PTask) -> None:
     if has_mark(task, "persist"):
         stateful_nodes: list[tuple[PTask | PNode, str | None]] = []
         for name in node_and_neighbors(session.dag, task.signature):
-            node = session.dag.nodes[name].value
+            node = session.dag.nodes[name]
             if isinstance(node, PProvisionalNode):
                 continue
             stateful_nodes.append((node, node.state()))

src/_pytask/profile.py

@@ -210,7 +210,7 @@ def pytask_profile_add_info_on_task(
             if successors:
                 sum_bytes = 0
                 for successor in successors:
-                    node = session.dag.nodes[successor].node_or_raise()
+                    node = session.dag.nodes[successor]
                     if isinstance(node, PPathNode):
                         with suppress(FileNotFoundError):
                             sum_bytes += node.path.stat().st_size

src/_pytask/shared.py

@@ -84,7 +84,7 @@ def reduce_names_of_multiple_nodes(
     """Reduce the names of multiple nodes in the DAG."""
     short_names = []
     for name in names:
-        node = dag.nodes[name].value
+        node = dag.nodes[name]
 
         if isinstance(node, PTask):
             short_name = format_task_name(node, editor_url_scheme="no_link").plain

src/_pytask/skipping.py

@@ -9,6 +9,7 @@
 from _pytask.mark import Mark
 from _pytask.mark_utils import get_marks
 from _pytask.mark_utils import has_mark
+from _pytask.node_protocols import PTask
 from _pytask.outcomes import Skipped
 from _pytask.outcomes import SkippedAncestorFailed
 from _pytask.outcomes import SkippedUnchanged
@@ -17,7 +18,6 @@
 from _pytask.provisional_utils import collect_provisional_products
 
 if TYPE_CHECKING:
-    from _pytask.node_protocols import PTask
     from _pytask.reports import ExecutionReport
     from _pytask.session import Session
 
@@ -97,9 +97,13 @@ def pytask_execute_task_process_report(
             report.outcome = TaskOutcome.SKIP
 
             for descending_task_name in descending_tasks(task.signature, session.dag):
-                descending_task = session.dag.nodes[
-                    descending_task_name
-                ].task_or_raise()
+                descending_task = session.dag.nodes[descending_task_name]
+                if not isinstance(descending_task, PTask):
+                    msg = (
+                        f"Expected descending task for signature "
+                        f"{descending_task_name!r}."
+                    )
+                    raise TypeError(msg)
                 descending_task.markers.append(
                     Mark(
                         "skip",