Simplify internal DAG types

Author: tobiasraabe

src/_pytask/dag_graph.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import itertools
+from dataclasses import dataclass
+from dataclasses import field
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import cast
@@ -17,56 +19,18 @@ class NoCycleError(Exception):
     """Raised when no cycle is found in a graph."""
 
 
-class NodeView:
-    """A minimal mapping-like view over node attributes."""
-
-    def __init__(self, node_attributes: dict[str, dict[str, Any]]) -> None:
-        self._node_attributes = node_attributes
-
-    def __getitem__(self, node: str) -> dict[str, Any]:
-        return self._node_attributes[node]
-
-    def __iter__(self) -> Iterator[str]:
-        return iter(self._node_attributes)
-
-    def __len__(self) -> int:
-        return len(self._node_attributes)
-
-    def __contains__(self, node: object) -> bool:
-        return node in self._node_attributes
-
-
-class UndirectedGraph:
-    """A minimal undirected graph used for validation tests."""
-
-    def __init__(
-        self,
-        node_attributes: dict[str, dict[str, Any]],
-        adjacency: dict[str, dict[str, None]],
-        graph_attributes: dict[str, Any],
-    ) -> None:
-        self._node_attributes = {
-            node: attributes.copy() for node, attributes in node_attributes.items()
-        }
-        self._adjacency = {
-            node: neighbors.copy() for node, neighbors in adjacency.items()
-        }
-        self.graph = graph_attributes.copy()
-        self.nodes = NodeView(self._node_attributes)
-
-    def is_directed(self) -> bool:
-        return False
-
-
+@dataclass
 class DiGraph:
     """A minimal directed graph tailored to pytask's needs."""
 
-    def __init__(self) -> None:
-        self._node_attributes: dict[str, dict[str, Any]] = {}
-        self._successors: dict[str, dict[str, None]] = {}
-        self._predecessors: dict[str, dict[str, None]] = {}
-        self.graph: dict[str, Any] = {}
-        self.nodes = NodeView(self._node_attributes)
+    _node_attributes: dict[str, dict[str, Any]] = field(default_factory=dict)
+    _successors: dict[str, dict[str, None]] = field(default_factory=dict)
+    _predecessors: dict[str, dict[str, None]] = field(default_factory=dict)
+    graph: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def nodes(self) -> dict[str, dict[str, Any]]:
+        return self._node_attributes
 
     def add_node(self, node_name: str, **attributes: Any) -> None:
         if node_name not in self._node_attributes:
@@ -138,16 +102,6 @@ def set_node_attributes(self, values: dict[str, Any], name: str) -> None:
             if node in self._node_attributes:
                 self._node_attributes[node][name] = value
 
-    def to_undirected(self) -> UndirectedGraph:
-        adjacency = {
-            node: {
-                **self._predecessors[node],
-                **self._successors[node],
-            }
-            for node in self._node_attributes
-        }
-        return UndirectedGraph(self._node_attributes, adjacency, self.graph)
-
     def to_networkx(self) -> Any:
         nx = cast("Any", import_optional_dependency("networkx"))
         graph = nx.DiGraph()

tests/test_dag_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from contextlib import ExitStack as does_not_raise  # noqa: N813
+from dataclasses import dataclass
 from pathlib import Path
 
 import pytest
@@ -143,10 +144,15 @@ def test_extract_priorities_from_tasks(tasks, expectation, expected):
         assert result == expected
 
 
-def test_raise_error_for_undirected_graphs(dag):
-    undirected_graph = dag.to_undirected()
+@dataclass
+class _UndirectedGraphStub:
+    def is_directed(self):
+        return False
+
+
+def test_raise_error_for_undirected_graphs():
     with pytest.raises(ValueError, match="Only directed graphs have a"):
-        TopologicalSorter.from_dag(undirected_graph)
+        TopologicalSorter.from_dag(_UndirectedGraphStub())  # type: ignore[arg-type]
 
 
 def test_raise_error_for_cycle_in_graph(dag):