diff --git a/src/daglib/dag.py b/src/daglib/dag.py index 0273280..9befa0b 100644 --- a/src/daglib/dag.py +++ b/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: