gabriel / musehub public
dag.py python
100 lines 3.6 KB
Raw
sha256:da49a05fd62cda46a7d73ec53a8d0adc5835d2070f6f0c51b12233a673c2e109 docs(mwp-1): mark Phase 5 complete, all acceptance criteria… Sonnet 4.6 26 days ago
1 """Core DAG primitives — nodes, edges, and the identity graph container."""
2
3 from collections import defaultdict
4 from collections.abc import Mapping
5 from enum import Enum
6 from typing import Iterable
7
8 class NodeType(str, Enum):
9 HUMAN = "human"
10 AGENT = "agent"
11 ORG = "org"
12
13 class EdgeType(str, Enum):
14 SPAWNS = "spawns"
15 MEMBER_OF = "member_of"
16
17 class GraphEdge:
18 __slots__ = ("from_handle", "to_handle", "edge_type")
19
20 def __init__(self, from_handle: str, to_handle: str, edge_type: EdgeType) -> None:
21 self.from_handle = from_handle
22 self.to_handle = to_handle
23 self.edge_type = edge_type
24
25 class IdentityDAG:
26 """Directed acyclic graph of identity nodes.
27
28 Edges are stored as an adjacency list: from_handle → list[to_handle].
29 Node types are optional metadata; from_edges() omits them.
30 """
31
32 def __init__(
33 self,
34 node_types: dict[str, NodeType],
35 adj: dict[str, list[str]],
36 radj: dict[str, list[str]],
37 ) -> None:
38 self._node_types: dict[str, NodeType] = node_types
39 self._adj: dict[str, list[str]] = adj # from → [to, ...]
40 self._radj: dict[str, list[str]] = radj # to → [from, ...]
41
42 # ── factories ─────────────────────────────────────────────────────────────
43
44 @classmethod
45 def empty(cls) -> "IdentityDAG":
46 return cls({}, defaultdict(list), defaultdict(list))
47
48 @classmethod
49 def from_edges(cls, edges: Iterable[GraphEdge]) -> "IdentityDAG":
50 adj: dict[str, list[str]] = defaultdict(list)
51 radj: dict[str, list[str]] = defaultdict(list)
52 for e in edges:
53 adj[e.from_handle].append(e.to_handle)
54 radj[e.to_handle].append(e.from_handle)
55 return cls({}, adj, radj)
56
57 @classmethod
58 def from_nodes_and_edges(
59 cls,
60 nodes: dict[str, NodeType],
61 edges: Iterable[GraphEdge],
62 ) -> "IdentityDAG":
63 adj: dict[str, list[str]] = defaultdict(list)
64 radj: dict[str, list[str]] = defaultdict(list)
65 for e in edges:
66 adj[e.from_handle].append(e.to_handle)
67 radj[e.to_handle].append(e.from_handle)
68 return cls(dict(nodes), adj, radj)
69
70 # ── mutation ───────────────────────────────────────────────────────────────
71
72 def add_node(self, handle: str, node_type: NodeType) -> None:
73 self._node_types[handle] = node_type
74
75 def add_edge(self, from_handle: str, to_handle: str) -> None:
76 self._adj[from_handle].append(to_handle)
77 self._radj[to_handle].append(from_handle)
78
79 def remove_edge(self, from_handle: str, to_handle: str) -> None:
80 self._adj[from_handle].remove(to_handle)
81 self._radj[to_handle].remove(from_handle)
82
83 # ── queries ────────────────────────────────────────────────────────────────
84
85 @property
86 def nodes(self) -> Mapping[str, NodeType]:
87 return self._node_types
88
89 def successors(self, handle: str) -> list[str]:
90 return list(self._adj.get(handle, []))
91
92 def predecessors(self, handle: str) -> list[str]:
93 return list(self._radj.get(handle, []))
94
95 def all_handles(self) -> set[str]:
96 handles: set[str] = set(self._node_types.keys())
97 for h, nbrs in self._adj.items():
98 handles.add(h)
99 handles.update(nbrs)
100 return handles
File History 1 commit
sha256:da49a05fd62cda46a7d73ec53a8d0adc5835d2070f6f0c51b12233a673c2e109 docs(mwp-1): mark Phase 5 complete, all acceptance criteria… Sonnet 4.6 26 days ago