98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
"""Topological ordering with parallel execution levels.
|
|
|
|
Kahn's algorithm, but instead of emitting one node at a time we emit
|
|
*levels*: each level contains every node whose in-degree has just hit zero,
|
|
so all nodes in a level can run in parallel. Deterministic tie-break by node
|
|
key (string) so output is reproducible. Raises CycleError listing the
|
|
remaining (unresolved) nodes if the graph has a cycle.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections import defaultdict, deque
|
|
from typing import Any, Iterable, List, Tuple
|
|
|
|
|
|
class CycleError(ValueError):
|
|
def __init__(self, remaining: List[Any]):
|
|
super().__init__(f"graph contains a cycle; unresolved nodes: {remaining!r}")
|
|
self.remaining = remaining
|
|
|
|
|
|
def topo_levels(
|
|
nodes: Iterable[Any],
|
|
edges: Iterable[Tuple[Any, Any]],
|
|
) -> List[List[Any]]:
|
|
"""Return levels of nodes in dependency order.
|
|
|
|
`edges` are (prerequisite, dependent) pairs: prerequisite must come before
|
|
dependent. Nodes with no edges still appear (in their own level 0 group).
|
|
"""
|
|
node_set = list(dict.fromkeys(nodes)) # preserve insertion order, dedupe
|
|
index = {n: i for i, n in enumerate(node_set)}
|
|
|
|
# validate edges reference known nodes
|
|
for a, b in edges:
|
|
if a not in index:
|
|
raise KeyError(f"edge references unknown node: {a!r}")
|
|
if b not in index:
|
|
raise KeyError(f"edge references unknown node: {b!r}")
|
|
|
|
indeg = {n: 0 for n in node_set}
|
|
adj = defaultdict(list)
|
|
for a, b in edges:
|
|
adj[a].append(b)
|
|
indeg[b] += 1
|
|
|
|
# deterministic: process ready queue sorted by key
|
|
ready = sorted([n for n in node_set if indeg[n] == 0], key=_key)
|
|
levels: List[List[Any]] = []
|
|
processed = 0
|
|
# Kahn in waves
|
|
current = ready
|
|
while current:
|
|
level = sorted(current, key=_key)
|
|
levels.append(level)
|
|
next_ready: List[Any] = []
|
|
for n in level:
|
|
processed += 1
|
|
for m in adj[n]:
|
|
indeg[m] -= 1
|
|
if indeg[m] == 0:
|
|
next_ready.append(m)
|
|
# dedupe next_ready (a node can be pushed twice if two prereqs resolve same wave? no, indeg hits 0 once)
|
|
current = next_ready
|
|
|
|
if processed != len(node_set):
|
|
remaining = [n for n in node_set if indeg[n] > 0]
|
|
raise CycleError(sorted(remaining, key=_key))
|
|
|
|
# flatten check: every node appears exactly once
|
|
seen = []
|
|
for lvl in levels:
|
|
seen.extend(lvl)
|
|
assert len(seen) == len(node_set) == len(set(seen)), "internal invariant violated"
|
|
return levels
|
|
|
|
|
|
def _key(n: Any):
|
|
return (type(n).__name__, n) if _comparable(n) else (type(n).__name__, repr(n))
|
|
|
|
|
|
def _comparable(n: Any) -> bool:
|
|
# Always comparable within the same concrete type for primitives we care about.
|
|
return isinstance(n, (int, float, str, bytes))
|
|
|
|
|
|
def flatten(levels: List[List[Any]]) -> List[Any]:
|
|
out = []
|
|
for lvl in levels:
|
|
out.extend(lvl)
|
|
return out
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# smoke
|
|
nodes = ["a", "b", "c", "d"]
|
|
edges = [("a", "c"), ("b", "c"), ("c", "d")]
|
|
print(topo_levels(nodes, edges))
|