brought out hooks

This commit is contained in:
John Lancaster
2026-02-21 00:46:35 -06:00
parent 1a2622cd2e
commit 78b00857c5

View File

@@ -1,7 +1,7 @@
from __future__ import annotations from __future__ import annotations
from collections import defaultdict 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 typing import Generic, TypeVar
from .set import DAGSet from .set import DAGSet
@@ -11,7 +11,7 @@ type dictset = defaultdict[T, set[T]]
T = TypeVar("T") T = TypeVar("T")
class DAG(Generic[T], MutableMapping[T, MutableSet[T]]): class DAG(Generic[T], MutableMapping[T, DAGSet[T]]):
""" """
DAG adjacency map: DAG adjacency map:
- g[u] -> mutable set-like view of successors of u - 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] _succ: dictset[T]
_pred: 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: def __init__(self) -> None:
self._succ: dictset[T] = defaultdict(set) self._succ: dictset[T] = defaultdict(set)
self._pred: dictset[T] = defaultdict(set) self._pred: dictset[T] = defaultdict(set)
@@ -34,21 +37,23 @@ class DAG(Generic[T], MutableMapping[T, MutableSet[T]]):
# --- MutableMapping interface --- # --- MutableMapping interface ---
def __getitem__(self, u: T) -> MutableSet[T]: def __getitem__(self, u: T) -> DAGSet[T]:
dagset = DAGSet(self._succ.get(u)) dagset = DAGSet(self._succ.get(u))
def on_add(v: T) -> None: def _on_add(v: T) -> None:
self._succ[u] |= {v} self._succ[u] |= {v}
self._pred[v] |= {u} 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._succ[u] -= {v}
self._pred[v] -= {u} 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_add = _on_add
dagset.on_remove = on_remove dagset.on_remove = _on_remove
return dagset return dagset
def __setitem__(self, u: T, vs: Iterable[T]) -> None: def __setitem__(self, u: T, vs: Iterable[T]) -> None: