Frontier walks, typed walks, parents, locations and subtree hashes in Rust

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "fast-walk"
-version = "0.2.2"
+version = "0.3.0"
 edition = "2024"
 license = "MIT"
 readme = "README.md"

README.md

@@ -37,6 +37,42 @@ for node in walk_unordered(tree):
   depth-first pre-order visitation. `ast.walk` does not document an order,
   so most callers can safely use `walk_unordered`.
 
+### Visiting only some node types
+
+`walk_frontier(node, kinds)` returns the descendants whose exact type is in
+`kinds`, in depth-first pre-order, without descending below a match. That is
+the set `ast.NodeVisitor.generic_visit` hands to `visit_<kind>` methods, so a
+visitor whose only methods are `visit_<kind>` can descend with:
+
+```python
+def generic_visit(self, node):
+    for child in walk_frontier(node, (ast.Call, ast.Attribute)):
+        self.visit(child)
+```
+
+`walk_frontier_edges` returns `(node, parent, field, index)` tuples instead
+(`index` is `-1` for a non-list field), which is what a `NodeTransformer`
+needs to put a replacement back.
+
+### Other helpers
+
+- `walk_of_types(node, kinds, prune=None)` — every node of the given exact
+  types, in pre-order, descending into matches but not below a `prune` type.
+- `walk_frontier_events(node, kinds, bracket)` — the frontier walk as
+  `(event, node)` pairs, with enter/leave events around `bracket` types, for a
+  visitor that does work around its generic descent.
+- `set_parents(node)` — `child.parent = node` for the whole tree.
+- `fix_missing_locations(node)` — `ast.fix_missing_locations` without
+  recursion.
+- `subtree_hashes(node)` — `{id(n): hash}` such that equal hashes mean equal
+  `ast.dump` output (leaf types included, so `1` and `True` differ).
+  `subtree_hash(node)` is the root's entry alone.
+
+These, and the frontier walks, read each child by its `_fields` name, so
+they also work on nodes a transformer built with keyword arguments in
+another order. `walk_dfs` and `walk_unordered` read fields by position,
+which is only guaranteed for parsed trees.
+
 ## Performance
 
 Benchmark on CPython 3.13, walking the AST of `difflib.py` (~2000 lines,

fast_walk.pyi

@@ -23,6 +23,71 @@ def walk_unordered(node: ast.AST) -> list[ast.AST]:
     25% faster than :func:`walk_dfs` on real Python source.
     """
 
+def walk_frontier(
+    node: ast.AST, kinds: tuple[type[ast.AST], ...]
+) -> list[ast.AST]:
+    """Return the descendants of `node` whose exact type is in `kinds`, in
+    depth-first pre-order, without descending below any match. `node`
+    itself is never included.
+
+    These are the nodes ``ast.NodeVisitor.generic_visit`` dispatches to a
+    ``visit_<kind>`` method when every other node falls through to the
+    generic descent.
+    """
+
+def walk_frontier_edges(
+    node: ast.AST, kinds: tuple[type[ast.AST], ...]
+) -> list[tuple[ast.AST, ast.AST, str, int]]:
+    """Like :func:`walk_frontier`, but each match comes as
+    ``(node, parent, field, index)``: it was read from ``parent.<field>``
+    when ``index`` is ``-1``, else from ``parent.<field>[index]``.
+    """
+
+def walk_of_types(
+    node: ast.AST,
+    kinds: tuple[type[ast.AST], ...],
+    prune: tuple[type[ast.AST], ...] | None = None,
+) -> list[ast.AST]:
+    """Return every node under `node` (`node` included) whose exact type is
+    in `kinds`, in depth-first pre-order. Unlike :func:`walk_frontier`,
+    matches are descended into, except a node whose exact type is in
+    `prune`: it is returned if it matches and its children are skipped.
+    `node` itself is always descended.
+    """
+
+def subtree_hash(node: ast.AST) -> int:
+    """``subtree_hashes(node)[id(node)]``, without building the dict."""
+
+def set_parents(node: ast.AST) -> None:
+    """Set ``child.parent`` for every node under `node`, visiting parents in
+    depth-first pre-order. A node reachable from two parents keeps the last.
+    """
+
+def walk_frontier_events(
+    node: ast.AST,
+    kinds: tuple[type[ast.AST], ...],
+    bracket: tuple[type[ast.AST], ...],
+) -> list[tuple[int, ast.AST]]:
+    """The calls a ``NodeVisitor`` makes descending from `node`.
+
+    ``(0, n)`` is a descendant whose exact type is in `kinds`, dispatched and
+    not descended. ``(1, n)`` and ``(2, n)`` surround the descent of a node
+    whose exact type is in `bracket`. `node` itself is always descended.
+    """
+
+def fix_missing_locations[T: ast.AST](node: T) -> T:
+    """``ast.fix_missing_locations`` without recursing once per tree level."""
+
+def subtree_hashes(node: ast.AST) -> dict[int, int]:
+    """A structural hash per node under `node`, keyed by ``id(node)``.
+
+    Subtrees with equal ``ast.dump`` output get equal hashes. A leaf value's
+    type is mixed in, so ``1``, ``1.0`` and ``True`` hash apart, and a value
+    that is not a str, bytes, int, bool, None or Ellipsis is hashed by its
+    ``repr``, so NaNs hash together, ``-0.0`` and ``0.0`` apart, and an
+    unhashable constant still hashes.
+    """
+
 def walk(node: ast.AST) -> list[ast.AST]:
     """Deprecated. Use :func:`walk_dfs` for explicit depth-first order or
     :func:`walk_unordered` for the faster order-agnostic variant.