walk.py
python
sha256:fd35b52294fe17892f05804659f29fa44e583018c702ea23073b154375506648
chore: bump version to 0.2.0.dev4 for nightly.4
Sonnet 5
patch
53 days ago
| 1 | """Generic async DAG walker for MuseHub. |
| 2 | |
| 3 | Mirrors ``muse.core.graph.walk_dag`` but accepts an async adjacency function, |
| 4 | making it suitable for SQLAlchemy AsyncSession queries and other async I/O. |
| 5 | |
| 6 | The canonical usage pattern is to write an async adjacency closure over a |
| 7 | session or an in-memory dict, then drive the walk with ``walk_dag_async``: |
| 8 | |
| 9 | async def adjacency(commit_id: str) -> list[str]: |
| 10 | row = await session.get(MusehubCommit, commit_id) |
| 11 | return row.parent_ids if row else [] |
| 12 | |
| 13 | async for cid in walk_dag_async(head, adjacency, max_nodes=10_000): |
| 14 | ... |
| 15 | """ |
| 16 | from __future__ import annotations |
| 17 | |
| 18 | import collections |
| 19 | from collections.abc import AsyncIterator, Callable, Coroutine, Iterable |
| 20 | from typing import Literal, TypeVar |
| 21 | |
| 22 | T = TypeVar("T", bound=object) |
| 23 | |
| 24 | AdjacencyFn = Callable[[T], Coroutine[None, None, list[T]]] |
| 25 | |
| 26 | |
| 27 | async def walk_dag_async( |
| 28 | starts: "T | list[T] | Iterable[T]", |
| 29 | adjacency: AdjacencyFn, |
| 30 | *, |
| 31 | order: "Literal['bfs', 'dfs']" = "bfs", |
| 32 | exclude: "set[T] | frozenset[T] | None" = None, |
| 33 | max_nodes: "int | None" = None, |
| 34 | ) -> "AsyncIterator[T]": |
| 35 | """Async BFS/DFS over any DAG with an async adjacency function. |
| 36 | |
| 37 | Parameters |
| 38 | ---------- |
| 39 | starts: |
| 40 | A single node, or a list/iterable of start nodes. Multi-root BFS |
| 41 | seeds the walk from all of them simultaneously. |
| 42 | adjacency: |
| 43 | Async callable: given a node, returns its immediate neighbours. |
| 44 | Called once per visited node, before the node is yielded. |
| 45 | order: |
| 46 | ``"bfs"`` (breadth-first, default) or ``"dfs"`` (depth-first). |
| 47 | exclude: |
| 48 | Nodes treated as already-visited. They and their subtrees are |
| 49 | never yielded. The caller's set is never mutated. |
| 50 | max_nodes: |
| 51 | Hard cap on the number of nodes yielded. ``None`` means no limit. |
| 52 | |
| 53 | Yields |
| 54 | ------ |
| 55 | T |
| 56 | Each reachable node at most once, in BFS or DFS visit order. |
| 57 | """ |
| 58 | if order not in ("bfs", "dfs"): |
| 59 | raise ValueError(f"order must be 'bfs' or 'dfs', got {order!r}") |
| 60 | |
| 61 | if isinstance(starts, (str, bytes)) or not hasattr(starts, "__iter__"): |
| 62 | start_list: list[T] = [starts] # type: ignore[list-item] |
| 63 | else: |
| 64 | start_list = list(starts) # type: ignore[arg-type] |
| 65 | |
| 66 | seen: set[T] = set(exclude) if exclude is not None else set() |
| 67 | yielded = 0 |
| 68 | queue: collections.deque = collections.deque(start_list) |
| 69 | |
| 70 | while queue: |
| 71 | if max_nodes is not None and yielded >= max_nodes: |
| 72 | return |
| 73 | |
| 74 | node = queue.popleft() if order == "bfs" else queue.pop() |
| 75 | |
| 76 | if node in seen: |
| 77 | continue |
| 78 | seen.add(node) |
| 79 | |
| 80 | neighbours = await adjacency(node) |
| 81 | |
| 82 | yield node |
| 83 | yielded += 1 |
| 84 | |
| 85 | for neighbour in neighbours: |
| 86 | if neighbour not in seen: |
| 87 | queue.append(neighbour) |
File History
1 commit
sha256:fd35b52294fe17892f05804659f29fa44e583018c702ea23073b154375506648
chore: bump version to 0.2.0.dev4 for nightly.4
Sonnet 5
patch
53 days ago