Weak Relationship Management

relationship.py

import gc
from operator import call
from weakref import ref, ReferenceType, WeakKeyDictionary
from typing import TypeVar, Generic, Set, List, Dict

T = TypeVar("T")

class WeakChildRelationship(Generic[T]):
    __slots__ = ("_dict",)

    def __init__(self) -> None:
        self._dict: WeakKeyDictionary[T, Set[ReferenceType[T]]] = WeakKeyDictionary()
        gc.callbacks.append(self._gc_cleanup)

    def add_child(self, parent: T, child: T) -> None:
        children = self._dict.get(parent)
        if children is None:
            children = self._dict[parent] = set()
        children.add(ref(child))

    def get_children(self, parent: T) -> List[T]:
        children = self._dict.get(parent)
        if children is None:
            return []
        return list(filter(None, map(call, children)))

    def _gc_cleanup(self, phase: str, info: Dict[str, int]) -> None:
        if phase == "stop" and info.get("generation") == 2:
            for key, children in list(self._dict.items()):
                live_children = list(filter(None, map(call, children)))
                if len(children) == len(live_children):
                    continue
                self._dict[key] = children.intersection(map(ref, live_children))


if __name__ == "__main__":
    RELATIONSHIP = WeakChildRelationship()

    class Node: pass

    n1 = Node()
    n2 = Node()
    RELATIONSHIP.add_child(n1, n2)
    assert RELATIONSHIP.get_children(n1) == [n2]
    del n2
    assert RELATIONSHIP.get_children(n1) == []

    assert len(RELATIONSHIP._dict[n1]) == 1
    gc.collect()
    assert len(RELATIONSHIP._dict[n1]) == 0