test hooks

src/daglib/dag.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from collections import defaultdict
 from collections.abc import Iterable, Iterator, MutableMapping, MutableSet
-from typing import DefaultDict, Generic, Set, TypeVar
+from typing import Generic, Set, TypeVar
 
 from .set import DAGSet
 
@@ -20,15 +20,27 @@ class DAG(Generic[T], MutableMapping[T, MutableSet[T]]):
     """
 
     def __init__(self) -> None:
-        self._succ: DefaultDict[T, Set[T]] = defaultdict(set)
-        self.reverse: DefaultDict[T, Set[T]] = defaultdict(set)
+        self._succ: defaultdict[T, Set[T]] = defaultdict(set)
+        self.reverse: defaultdict[T, Set[T]] = defaultdict(set)
 
     # --- MutableMapping interface ---
 
     def __getitem__(self, u: T) -> MutableSet[T]:
-        # defaultdict semantics: touch key
-        _ = self._succ[u]
-        return DAGSet(self, u)
+        dagset = DAGSet(self._succ.get(u))
+
+        def on_add(v: T) -> None:
+            self._succ[u].add(v)
+            self.reverse[v].add(u)
+            print(f"Adding edge {u} -> {v}")
+
+        def on_remove(v: T) -> None:
+            self._succ[u].remove(v)
+            self.reverse[v].remove(u)
+            print(f"Removing edge {u} -> {v}")
+
+        dagset.on_add = on_add
+        dagset.on_remove = on_remove
+        return dagset
 
     def __setitem__(self, u: T, vs: Iterable[T]) -> None:
         view = DAGSet(self, u)
@@ -45,7 +57,7 @@ class DAG(Generic[T], MutableMapping[T, MutableSet[T]]):
         return len(self._succ)
 
     def __repr__(self) -> str:
-        return f"DagAdj({dict(self._succ)!r})"
+        return f"{self.__class__.__name__}{dict(self._succ)!r}"
 
     # --- core edge/node operations (single source of truth) ---
 
@@ -53,6 +65,11 @@ class DAG(Generic[T], MutableMapping[T, MutableSet[T]]):
         if u == v:
             raise ValueError("Self-loops are not allowed in a DAG")
 
+        if u not in self._succ:
+            self._succ[u] = set()
+        if v not in self.reverse:
+            self.reverse[v] = set()
+
         if v not in self._succ[u]:
             self._succ[u].add(v)
             self.reverse[v].add(u)