brought out hooks

src/daglib/dag.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from collections.abc import Iterable, Iterator, MutableMapping, MutableSet
+from collections.abc import Callable, Iterable, Iterator, MutableMapping
 from typing import Generic, TypeVar
 
 from .set import DAGSet
@@ -11,7 +11,7 @@ type dictset = defaultdict[T, set[T]]
 T = TypeVar("T")
 
 
-class DAG(Generic[T], MutableMapping[T, MutableSet[T]]):
+class DAG(Generic[T], MutableMapping[T, DAGSet[T]]):
     """
     DAG adjacency map:
       - g[u] -> mutable set-like view of successors of u
@@ -24,6 +24,9 @@ class DAG(Generic[T], MutableMapping[T, MutableSet[T]]):
     _succ: dictset[T]
     _pred: dictset[T]
 
+    on_add: Callable[[T, T], None] | None = None
+    on_remove: Callable[[T, T], None] | None = None
+
     def __init__(self) -> None:
         self._succ: dictset[T] = defaultdict(set)
         self._pred: dictset[T] = defaultdict(set)
@@ -34,21 +37,23 @@ class DAG(Generic[T], MutableMapping[T, MutableSet[T]]):
 
     # --- MutableMapping interface ---
 
-    def __getitem__(self, u: T) -> MutableSet[T]:
+    def __getitem__(self, u: T) -> DAGSet[T]:
         dagset = DAGSet(self._succ.get(u))
 
-        def on_add(v: T) -> None:
+        def _on_add(v: T) -> None:
             self._succ[u] |= {v}
             self._pred[v] |= {u}
-            print(f"Adding edge {u} -> {v}")
+            if self.on_add is not None:
+                self.on_add(u, v)
 
-        def on_remove(v: T) -> None:
+        def _on_remove(v: T) -> None:
             self._succ[u] -= {v}
             self._pred[v] -= {u}
-            print(f"Removing edge {u} -> {v}")
+            if self.on_remove is not None:
+                self.on_remove(u, v)
 
-        dagset.on_add = on_add
-        dagset.on_remove = on_remove
+        dagset.on_add = _on_add
+        dagset.on_remove = _on_remove
         return dagset
 
     def __setitem__(self, u: T, vs: Iterable[T]) -> None: