diff --git a/Cargo.lock b/Cargo.lock index 429eb88..2dd2439 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,7 +4,7 @@ version = 4 [[package]] name = "fast-walk" -version = "0.2.2" +version = "0.3.0" dependencies = [ "pyo3", ] diff --git a/Cargo.toml b/Cargo.toml index 6b16a40..245c66a 100644 --- a/Cargo.toml +++ b/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" diff --git a/README.md b/README.md index d0f5f96..e98d8a8 100644 --- a/README.md +++ b/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_` methods, so a +visitor whose only methods are `visit_` 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, diff --git a/fast_walk.pyi b/fast_walk.pyi index 8dd5e38..ebeebe7 100644 --- a/fast_walk.pyi +++ b/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_`` 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.`` + when ``index`` is ``-1``, else from ``parent.[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. diff --git a/src/fields.rs b/src/fields.rs new file mode 100644 index 0000000..aa4d449 --- /dev/null +++ b/src/fields.rs @@ -0,0 +1,676 @@ +//! Traversals that read a node's children by its type's `_fields` names. +//! +//! `walk_dfs` reads the first `len(_fields)` entries of the instance dict, +//! which holds for a parsed tree. A transformer can build a node whose dict +//! has another order (`ast.Call(lineno=1, func=f)` stores `lineno` first), +//! and these functions run on such trees, so each field is looked up by +//! name instead. +//! +//! A field read can run Python code (a property, a `__getattr__`, a value's +//! `__hash__` or `__repr__`) that drops the tree's own references, so every +//! node a walk will still touch is held as a `Strong`. + +use std::cell::RefCell; +use std::collections::HashMap; +use std::hash::{BuildHasherDefault, Hasher}; + +use pyo3::ffi::{self, PyObject, PyTypeObject}; +use pyo3::prelude::*; +use pyo3::types::{PyList, PyTuple, PyType}; +use pyo3::{PyTypeInfo, intern}; + +use crate::{Edge, get_instance_dict_fast, issubclass_of_ast, resolve_base_types}; + +/// An owned reference, dropped with the GIL held. +pub(crate) struct Strong(*mut PyObject); + +impl Strong { + /// # Safety + /// `object` is a live object and the GIL is held. + #[inline(always)] + pub(crate) unsafe fn borrowed(object: *mut PyObject) -> Self { + unsafe { ffi::Py_INCREF(object) }; + Self(object) + } + + /// # Safety + /// `object` is a new reference the caller gives up. + #[inline(always)] + unsafe fn owned(object: *mut PyObject) -> Self { + Self(object) + } + + #[inline(always)] + pub(crate) fn as_ptr(&self) -> *mut PyObject { + self.0 + } + + #[inline(always)] + pub(crate) fn into_ptr(self) -> *mut PyObject { + let object = self.0; + std::mem::forget(self); + object + } +} + +impl Clone for Strong { + #[inline(always)] + fn clone(&self) -> Self { + unsafe { Self::borrowed(self.0) } + } +} + +impl Drop for Strong { + #[inline(always)] + fn drop(&mut self) { + unsafe { ffi::Py_DECREF(self.0) }; + } +} + +/// A list holding `items`, whose references it takes over. +pub(crate) fn into_pylist<'py>( + py: Python<'py>, + items: impl ExactSizeIterator, +) -> PyResult> { + unsafe { + let list = + Bound::from_owned_ptr_or_err(py, ffi::PyList_New(items.len() as ffi::Py_ssize_t))?; + let slots = (*(list.as_ptr() as *mut ffi::PyListObject)).ob_item; + for (position, item) in items.enumerate() { + *slots.add(position) = item.into_ptr(); + } + Ok(list) + } +} + +#[derive(Default)] +struct PointerHasher(u64); + +impl Hasher for PointerHasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, bytes: &[u8]) { + for &byte in bytes { + self.0 = (self.0 << 8 | byte as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15); + } + } + + fn write_usize(&mut self, value: usize) { + self.0 = ((value as u64) >> 4).wrapping_mul(0x9E37_79B9_7F4A_7C15); + } +} + +type PointerMap = HashMap>; + +const LINENO: u8 = 1; +const COL_OFFSET: u8 = 2; +const END_LINENO: u8 = 4; +const END_COL_OFFSET: u8 = 8; + +/// An AST type's `_fields` names, and which location `_attributes` it has. +/// It holds the type, so the address it is cached under cannot be reused by +/// another type. +struct TypeInfo { + _kind: Py, + fields: Vec>, + locations: u8, +} + +thread_local! { + static TYPE_INFO: RefCell>> = RefCell::new(PointerMap::default()); +} + +fn names_of(py: Python<'_>, kind: &Bound<'_, PyAny>, attribute: &str) -> Vec> { + let Ok(names) = kind.getattr(attribute) else { + return Vec::new(); + }; + let Ok(names) = names.cast_into::() else { + return Vec::new(); + }; + names + .iter() + .filter_map(|name| { + let name = name.cast_into::().ok()?; + let interned = pyo3::types::PyString::intern(py, name.to_str().ok()?); + Some(interned.into_any().unbind()) + }) + .collect() +} + +/// The cached `TypeInfo` for `kind`. Entries are boxed and never removed, +/// so the pointer stays valid for the life of the thread. A miss reads the +/// type's attributes, which can run Python code that walks a tree, so no +/// borrow of the cache is held across it. +fn type_info(py: Python<'_>, kind: *mut PyTypeObject) -> *const TypeInfo { + let cached = TYPE_INFO.with(|cache| { + cache + .borrow() + .get(&(kind as usize)) + .map(|entry| &**entry as *const TypeInfo) + }); + if let Some(info) = cached { + return info; + } + let kind_object = unsafe { Bound::from_borrowed_ptr(py, kind.cast::()) }; + let mut locations = 0; + for name in names_of(py, &kind_object, "_attributes") { + locations |= match name.bind(py).extract::<&str>().unwrap_or_default() { + "lineno" => LINENO, + "col_offset" => COL_OFFSET, + "end_lineno" => END_LINENO, + "end_col_offset" => END_COL_OFFSET, + _ => 0, + }; + } + let info = Box::new(TypeInfo { + fields: names_of(py, &kind_object, "_fields"), + locations, + _kind: kind_object.unbind(), + }); + TYPE_INFO + .with(|cache| &**cache.borrow_mut().entry(kind as usize).or_insert(info) as *const TypeInfo) +} + +/// Call `visit(field_index, value)` for each field `node` has, in `_fields` +/// order, as `ast.iter_fields` reads them: from the node's dict, else by +/// attribute lookup, which finds the class default of an optional field a +/// constructor left out (`ast.Constant(value=1)` has no `kind` entry). +/// `node` must be kept alive by the caller. The value is borrowed: a `visit` +/// that runs Python code holds it first. +unsafe fn for_each_field( + py: Python<'_>, + node: *mut PyObject, + info: &TypeInfo, + mut visit: impl FnMut(usize, *mut PyObject) -> PyResult<()>, +) -> PyResult<()> { + let dict = get_instance_dict_fast(node).map(|dict| unsafe { Strong::borrowed(dict) }); + for (index, name) in info.fields.iter().enumerate() { + if let Some(dict) = &dict { + let value = unsafe { ffi::PyDict_GetItemWithError(dict.as_ptr(), name.as_ptr()) }; + if !value.is_null() { + visit(index, value)?; + continue; + } + if unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + } + let inherited = unsafe { ffi::PyObject_GetAttr(node, name.as_ptr()) }; + if inherited.is_null() { + let error = PyErr::fetch(py); + if error.is_instance_of::(py) { + continue; + } + return Err(error); + } + let inherited = unsafe { Strong::owned(inherited) }; + visit(index, inherited.as_ptr())?; + } + Ok(()) +} + +/// The items of `list`, borrowed, read one at a time so Python code run +/// between them cannot leave the loop reading a resized list. +#[inline(always)] +unsafe fn for_each_item( + list: *mut PyObject, + mut visit: impl FnMut(ffi::Py_ssize_t, *mut PyObject) -> PyResult<()>, +) -> PyResult<()> { + let mut index = 0; + while index < unsafe { ffi::PyList_GET_SIZE(list) } { + visit(index, unsafe { ffi::PyList_GET_ITEM(list, index) })?; + index += 1; + } + Ok(()) +} + +struct Walk<'py> { + py: Python<'py>, + base: (*mut PyTypeObject, *mut PyTypeObject), + list_type: *mut PyTypeObject, +} + +impl<'py> Walk<'py> { + fn new(py: Python<'py>) -> PyResult { + Ok(Self { + py, + base: resolve_base_types(py)?, + list_type: PyList::type_object_raw(py), + }) + } + + /// Whether `value` is an `ast.AST` instance, at any depth of subclass. + #[inline(always)] + fn is_ast(&self, value: *mut PyObject) -> bool { + let kind = unsafe { ffi::Py_TYPE(value) }; + if issubclass_of_ast(kind, self.base) { + return true; + } + let first = unsafe { (*kind).tp_base }; + if first.is_null() || first == &raw mut ffi::PyBaseObject_Type { + return false; + } + unsafe { ffi::PyType_IsSubtype(kind, self.base.0) != 0 } + } + + /// Push the AST children of `node` onto `stack` so they pop in + /// `ast.iter_child_nodes` order. + fn push_children(&self, node: &Strong, stack: &mut Vec) -> PyResult<()> { + let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(node.as_ptr())) }; + let start = stack.len(); + unsafe { + for_each_field(self.py, node.as_ptr(), info, |_, value| { + if ffi::Py_TYPE(value) == self.list_type { + for_each_item(value, |_, item| { + if self.is_ast(item) { + stack.push(Strong::borrowed(item)); + } + Ok(()) + })?; + } else if self.is_ast(value) { + stack.push(Strong::borrowed(value)); + } + Ok(()) + })?; + } + stack[start..].reverse(); + Ok(()) + } + + /// Every node under `root`, `root` included, in depth-first pre-order. + fn pre_order(&self, root: &Bound<'_, PyAny>) -> PyResult> { + let mut result = Vec::new(); + let mut stack = vec![unsafe { Strong::borrowed(root.as_ptr()) }]; + while let Some(node) = stack.pop() { + self.push_children(&node, &mut stack)?; + result.push(node); + } + Ok(result) + } +} + +/// A set of exact types, sorted for a binary search. +pub(crate) struct TypeSet(Vec); + +impl TypeSet { + pub(crate) fn new(kinds: &Bound<'_, PyTuple>) -> PyResult { + let mut pointers = kinds + .iter() + .map(|kind| Ok(kind.cast_into::()?.as_type_ptr() as usize)) + .collect::>>()?; + pointers.sort_unstable(); + pointers.dedup(); + Ok(Self(pointers)) + } + + #[inline] + pub(crate) fn contains(&self, node: &Strong) -> bool { + self.0 + .binary_search(&(unsafe { ffi::Py_TYPE(node.as_ptr()) } as usize)) + .is_ok() + } +} + +/// Return every node under `node` (`node` included) whose exact type is in +/// `kinds`, in depth-first pre-order. Unlike `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. +#[pyfunction] +#[pyo3(signature = (node, kinds, prune=None))] +pub fn walk_of_types<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, + prune: Option>, +) -> PyResult> { + let kinds = TypeSet::new(&kinds)?; + let prune = prune.map(|prune| TypeSet::new(&prune)).transpose()?; + let walk = Walk::new(py)?; + let root = node.as_ptr(); + let mut result = Vec::new(); + let mut stack = vec![unsafe { Strong::borrowed(root) }]; + while let Some(current) = stack.pop() { + let pruned = current.as_ptr() != root + && prune.as_ref().is_some_and(|prune| prune.contains(¤t)); + if !pruned { + walk.push_children(¤t, &mut stack)?; + } + if kinds.contains(¤t) { + result.push(current); + } + } + into_pylist(py, result.into_iter()) +} + +/// Set `child.parent = parent` for every node under `node`, visiting parents +/// in depth-first pre-order and each parent's children in +/// `ast.iter_child_nodes` order. `node` itself is left alone. +#[pyfunction] +pub fn set_parents(py: Python<'_>, node: Bound<'_, PyAny>) -> PyResult<()> { + let walk = Walk::new(py)?; + let name = intern!(py, "parent"); + let mut children = Vec::new(); + let mut stack = vec![unsafe { Strong::borrowed(node.as_ptr()) }]; + while let Some(parent) = stack.pop() { + walk.push_children(&parent, &mut children)?; + for child in children.iter().rev() { + if unsafe { ffi::PyObject_SetAttr(child.as_ptr(), name.as_ptr(), parent.as_ptr()) } + == -1 + { + return Err(PyErr::fetch(py)); + } + } + stack.append(&mut children); + } + Ok(()) +} + +const VISIT: i64 = 0; +const ENTER: i64 = 1; +const LEAVE: i64 = 2; + +enum Step { + Child(Strong), + Leave(Strong), +} + +/// The calls a `NodeVisitor` makes descending from `node`, as a list of +/// `(event, node)` pairs: `(0, n)` for a descendant whose exact type is in +/// `kinds` (dispatched, not descended), and `(1, n)` / `(2, n)` around the +/// descent of a node whose exact type is in `bracket`. `node` itself is +/// always descended, and bracketed if its type is in `bracket`. +#[pyfunction] +pub fn walk_frontier_events<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, + bracket: Bound<'py, PyTuple>, +) -> PyResult> { + let kinds = TypeSet::new(&kinds)?; + let bracket = TypeSet::new(&bracket)?; + let walk = Walk::new(py)?; + let mut events: Vec<(i64, Strong)> = Vec::new(); + let mut stack: Vec = Vec::new(); + let mut children = Vec::new(); + let mut open = + |current: Strong, stack: &mut Vec, events: &mut Vec<(i64, Strong)>| -> PyResult<()> { + children.clear(); + walk.push_children(¤t, &mut children)?; + if bracket.contains(¤t) { + events.push((ENTER, current.clone())); + stack.push(Step::Leave(current)); + } + stack.extend(children.drain(..).map(Step::Child)); + Ok(()) + }; + open( + unsafe { Strong::borrowed(node.as_ptr()) }, + &mut stack, + &mut events, + )?; + while let Some(step) = stack.pop() { + match step { + Step::Leave(current) => events.push((LEAVE, current)), + Step::Child(current) if kinds.contains(¤t) => events.push((VISIT, current)), + Step::Child(current) => open(current, &mut stack, &mut events)?, + } + } + let codes = + [VISIT, ENTER, LEAVE].map(|code| code.into_pyobject(py).unwrap().into_any().unbind()); + unsafe { + let list = + Bound::from_owned_ptr_or_err(py, ffi::PyList_New(events.len() as ffi::Py_ssize_t))?; + let slots = (*(list.as_ptr() as *mut ffi::PyListObject)).ob_item; + for (position, (event, current)) in events.into_iter().enumerate() { + let pair = Bound::from_owned_ptr_or_err(py, ffi::PyTuple_New(2))?; + ffi::PyTuple_SET_ITEM( + pair.as_ptr(), + 0, + codes[event as usize].clone_ref(py).into_ptr(), + ); + ffi::PyTuple_SET_ITEM(pair.as_ptr(), 1, current.into_ptr()); + *slots.add(position) = pair.into_ptr(); + } + Ok(list) + } +} + +/// `ast.fix_missing_locations`, without recursing per tree level. Returns +/// `node`. +#[pyfunction] +pub fn fix_missing_locations<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, +) -> PyResult> { + let walk = Walk::new(py)?; + let names = [ + (LINENO, intern!(py, "lineno"), false), + (COL_OFFSET, intern!(py, "col_offset"), false), + (END_LINENO, intern!(py, "end_lineno"), true), + (END_COL_OFFSET, intern!(py, "end_col_offset"), true), + ]; + let one = 1i64.into_pyobject(py)?.into_any().unbind(); + let zero = 0i64.into_pyobject(py)?.into_any().unbind(); + let mut locations: Vec<[Py; 4]> = + vec![[one.clone_ref(py), zero.clone_ref(py), one, zero]]; + let mut stack: Vec<(Strong, usize)> = vec![(unsafe { Strong::borrowed(node.as_ptr()) }, 0)]; + let mut children = Vec::new(); + while let Some((current, inherited)) = stack.pop() { + let info = unsafe { &*type_info(py, ffi::Py_TYPE(current.as_ptr())) }; + let mut own = inherited; + if info.locations != 0 { + let dict = get_instance_dict_fast(current.as_ptr()) + .map(|dict| unsafe { Strong::borrowed(dict) }); + let mut values: [Py; 4] = + locations[inherited].each_ref().map(|v| v.clone_ref(py)); + for (slot, (flag, name, none_is_missing)) in names.iter().enumerate() { + if info.locations & flag == 0 { + continue; + } + let found = match &dict { + Some(dict) => unsafe { + ffi::PyDict_GetItemWithError(dict.as_ptr(), name.as_ptr()) + }, + None => std::ptr::null_mut(), + }; + if found.is_null() && unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + let missing = + found.is_null() || (*none_is_missing && unsafe { found == ffi::Py_None() }); + if missing { + unsafe { Bound::from_borrowed_ptr(py, current.as_ptr()) } + .setattr(name, values[slot].bind(py))?; + } else { + values[slot] = unsafe { Bound::from_borrowed_ptr(py, found) }.unbind(); + } + } + locations.push(values); + own = locations.len() - 1; + } + children.clear(); + walk.push_children(¤t, &mut children)?; + stack.extend(children.drain(..).map(|child| (child, own))); + } + Ok(node) +} + +#[inline(always)] +fn mix(hash: u64, value: u64) -> u64 { + (hash.rotate_left(5) ^ value).wrapping_mul(0x9E37_79B9_7F4A_7C15) +} + +/// A structural hash per node, keyed by `id(node)`. Two subtrees that would +/// give the same `ast.dump` get the same hash; a leaf value's type is mixed +/// in, so `1`, `1.0` and `True` (equal hashes in Python) stay apart. A leaf +/// that is not a str, bytes, int, bool, None or Ellipsis is hashed by its +/// `repr`, as `ast.dump` prints it: NaNs compare equal there and `-0.0` and +/// `0.0` do not, and a list constant has a repr where it has no hash. +#[pyfunction] +pub fn subtree_hashes<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, +) -> PyResult> { + let hashes = hash_subtrees(py, &node)?; + let result = pyo3::types::PyDict::new(py); + for (&pointer, &hash) in &hashes { + let key = unsafe { + Bound::from_owned_ptr_or_err( + py, + ffi::PyLong_FromVoidPtr(pointer as *mut std::ffi::c_void), + )? + }; + result.set_item(key, hash as i64)?; + } + Ok(result.into_any()) +} + +/// `subtree_hashes(node)[id(node)]`, without building the dict. +#[pyfunction] +pub fn subtree_hash(py: Python<'_>, node: Bound<'_, PyAny>) -> PyResult { + let hashes = hash_subtrees(py, &node)?; + Ok(hashed(&hashes, node.as_ptr())? as i64) +} + +fn hashed(hashes: &PointerMap, node: *mut PyObject) -> PyResult { + hashes.get(&(node as usize)).copied().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err("the tree changed while it was being hashed") + }) +} + +/// Whether `value`'s hash is what `ast.dump` compares: its type's hash is +/// C code that agrees with its repr. +#[inline(always)] +fn hashes_as_dumped(value: *mut PyObject) -> bool { + unsafe { + let kind = ffi::Py_TYPE(value); + kind == &raw mut ffi::PyUnicode_Type + || kind == &raw mut ffi::PyLong_Type + || kind == &raw mut ffi::PyBool_Type + || kind == &raw mut ffi::PyBytes_Type + || value == ffi::Py_None() + || value == ffi::Py_Ellipsis() + } +} + +fn hash_subtrees(py: Python<'_>, node: &Bound<'_, PyAny>) -> PyResult> { + let walk = Walk::new(py)?; + let nodes = walk.pre_order(node)?; + let mut hashes: PointerMap = PointerMap::default(); + hashes.reserve(nodes.len()); + let leaf = |value: *mut PyObject, hashes: &PointerMap| -> PyResult { + if walk.is_ast(value) { + if let Some(&hash) = hashes.get(&(value as usize)) { + return Ok(hash); + } + // A field that computes its node gives a new one on every read. + let child = unsafe { Bound::from_borrowed_ptr(py, value) }; + return hashed(&hash_subtrees(py, &child)?, value); + } + let kind = unsafe { ffi::Py_TYPE(value) } as u64; + let hash = unsafe { + if hashes_as_dumped(value) { + ffi::PyObject_Hash(value) + } else { + let repr = Bound::from_owned_ptr_or_err(py, ffi::PyObject_Repr(value))?; + ffi::PyObject_Hash(repr.as_ptr()) + } + }; + if hash == -1 && unsafe { !ffi::PyErr_Occurred().is_null() } { + return Err(PyErr::fetch(py)); + } + Ok(mix(kind, hash as u64)) + }; + for current in nodes.iter().rev() { + if hashes.contains_key(&(current.as_ptr() as usize)) { + continue; + } + let kind = unsafe { ffi::Py_TYPE(current.as_ptr()) }; + let info = unsafe { &*type_info(py, kind) }; + let mut hash = mix(0, kind as u64); + unsafe { + for_each_field(py, current.as_ptr(), info, |index, value| { + let value = Strong::borrowed(value); + let value_hash = if ffi::Py_TYPE(value.as_ptr()) == walk.list_type { + let mut items = mix(u64::MAX, ffi::PyList_GET_SIZE(value.as_ptr()) as u64); + for_each_item(value.as_ptr(), |_, item| { + let item = Strong::borrowed(item); + items = mix(items, leaf(item.as_ptr(), &hashes)?); + Ok(()) + })?; + items + } else { + leaf(value.as_ptr(), &hashes)? + }; + hash = mix(mix(hash, index as u64), value_hash); + Ok(()) + })?; + } + hashes.insert(current.as_ptr() as usize, hash); + } + Ok(hashes) +} + +impl Walk<'_> { + /// Push the AST children of `parent` onto `stack` so they pop in + /// `ast.iter_child_nodes` order, each with the slot it was read from. + fn push_child_edges(&self, parent: &Strong, stack: &mut Vec) -> PyResult<()> { + let info = unsafe { &*type_info(self.py, ffi::Py_TYPE(parent.as_ptr())) }; + let start = stack.len(); + unsafe { + for_each_field(self.py, parent.as_ptr(), info, |index, value| { + let key = info.fields[index].as_ptr(); + if ffi::Py_TYPE(value) == self.list_type { + for_each_item(value, |item, child| { + if self.is_ast(child) { + stack.push(Edge { + node: Strong::borrowed(child), + parent: parent.clone(), + key, + index: item, + }); + } + Ok(()) + })?; + } else if self.is_ast(value) { + stack.push(Edge { + node: Strong::borrowed(value), + parent: parent.clone(), + key, + index: -1, + }); + } + Ok(()) + })?; + } + stack[start..].reverse(); + Ok(()) + } +} + +/// Descendants of `node` (excluding `node`) whose exact type is in `kinds`, +/// in depth-first pre-order, without descending below any of them. This is +/// the set `ast.NodeVisitor.generic_visit` dispatches to `visit_` +/// methods when every other node falls through to the generic descent. +pub(crate) fn frontier_edges( + py: Python<'_>, + node: &Bound<'_, PyAny>, + kinds: &Bound<'_, PyTuple>, +) -> PyResult> { + let kinds = TypeSet::new(kinds)?; + let walk = Walk::new(py)?; + let mut result = Vec::new(); + let mut stack = Vec::new(); + walk.push_child_edges(&unsafe { Strong::borrowed(node.as_ptr()) }, &mut stack)?; + while let Some(edge) = stack.pop() { + if kinds.contains(&edge.node) { + result.push(edge); + } else { + walk.push_child_edges(&edge.node, &mut stack)?; + } + } + Ok(result) +} diff --git a/src/lib.rs b/src/lib.rs index 9837ae8..731968e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,6 +10,7 @@ //! `ast.walk` makes no ordering guarantee, so `walk_unordered` is a drop-in //! replacement wherever order doesn't matter. +mod fields; mod pydict; use std::cell::{Cell, RefCell}; @@ -17,7 +18,9 @@ use std::sync::atomic::{AtomicBool, Ordering}; use pyo3::exceptions::PyDeprecationWarning; use pyo3::ffi::{self, PyListObject, PyObject, PyTypeObject}; -use pyo3::types::{PyList, PyModule, PyType}; +use pyo3::types::{PyList, PyModule, PyTuple, PyType}; + +use fields::{Strong, frontier_edges, into_pylist}; use pyo3::{PyTypeInfo, prelude::*}; /// Open-addressed, direct-mapped lookup from `*mut PyTypeObject` to an @@ -514,6 +517,66 @@ fn walk_unordered<'py>(py: Python<'py>, node: Bound<'py, PyAny>) -> PyResult` when `index < 0`, else `parent.[index]`. +struct Edge { + node: Strong, + parent: Strong, + key: *mut PyObject, + index: ffi::Py_ssize_t, +} + +/// Return the descendants of `node` whose exact type is one of `kinds`, in +/// depth-first pre-order, without descending below any match. `node` +/// itself is never included. +#[pyfunction] +fn walk_frontier<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, +) -> PyResult> { + let edges = frontier_edges(py, &node, &kinds)?; + into_pylist(py, edges.into_iter().map(|edge| edge.node)) +} + +/// Like `walk_frontier`, but each match comes as `(node, parent, field, +/// index)`: the node was read from `parent.` when `index` is `-1`, +/// else from `parent.[index]`. +#[pyfunction] +fn walk_frontier_edges<'py>( + py: Python<'py>, + node: Bound<'py, PyAny>, + kinds: Bound<'py, PyTuple>, +) -> PyResult> { + let edges = frontier_edges(py, &node, &kinds)?; + unsafe { + let list_ptr = ffi::PyList_New(edges.len() as ffi::Py_ssize_t); + if list_ptr.is_null() { + return Err(PyErr::fetch(py)); + } + let list = Bound::from_owned_ptr(py, list_ptr); + let ob_item = (*(list_ptr as *mut ffi::PyListObject)).ob_item; + for (position, edge) in edges.into_iter().enumerate() { + let index = ffi::PyLong_FromSsize_t(edge.index); + if index.is_null() { + return Err(PyErr::fetch(py)); + } + let item = ffi::PyTuple_New(4); + if item.is_null() { + ffi::Py_DECREF(index); + return Err(PyErr::fetch(py)); + } + ffi::Py_INCREF(edge.key); + ffi::PyTuple_SET_ITEM(item, 0, edge.node.into_ptr()); + ffi::PyTuple_SET_ITEM(item, 1, edge.parent.into_ptr()); + ffi::PyTuple_SET_ITEM(item, 2, edge.key); + ffi::PyTuple_SET_ITEM(item, 3, index); + *ob_item.add(position) = item; + } + Ok(list) + } +} + static DEPRECATED_WALK_WARNED: AtomicBool = AtomicBool::new(false); /// Deprecated. Use `walk_dfs` for explicit depth-first order or @@ -556,6 +619,14 @@ fn fast_walk(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(walk, m)?)?; m.add_function(wrap_pyfunction!(walk_dfs, m)?)?; m.add_function(wrap_pyfunction!(walk_unordered, m)?)?; + m.add_function(wrap_pyfunction!(walk_frontier, m)?)?; + m.add_function(wrap_pyfunction!(walk_frontier_edges, m)?)?; + m.add_function(wrap_pyfunction!(fields::walk_of_types, m)?)?; + m.add_function(wrap_pyfunction!(fields::fix_missing_locations, m)?)?; + m.add_function(wrap_pyfunction!(fields::subtree_hashes, m)?)?; + m.add_function(wrap_pyfunction!(fields::subtree_hash, m)?)?; + m.add_function(wrap_pyfunction!(fields::set_parents, m)?)?; + m.add_function(wrap_pyfunction!(fields::walk_frontier_events, m)?)?; m.add_function(wrap_pyfunction!(_walk_count, m)?)?; Ok(()) } diff --git a/tests/test_fields.py b/tests/test_fields.py new file mode 100644 index 0000000..e732204 --- /dev/null +++ b/tests/test_fields.py @@ -0,0 +1,378 @@ +"""The functions here read each +node's children by `_fields` name, so they must agree with the stdlib on +parsed trees and on nodes a transformer built in another attribute order.""" + +from __future__ import annotations + +import ast +import copy +import gc +import sys +import textwrap + +import pytest + +from fast_walk import ( + fix_missing_locations, + set_parents, + subtree_hash, + subtree_hashes, + walk_frontier_events, + walk_of_types, +) + +SOURCE = textwrap.dedent(""" + import os + from x import y as z + def f(a: int = 1, *b, c=2.0, **d) -> str: + e = [g(h) for h in i if h] + if e and not a: + return k.l.m(n, *o, p=q) + with r() as s, t: + del u[v:w] + class C(D, metaclass=E): + x: int = True + async def m(self): + async for a in b: + await c + match a: + case [1, *rest] | [2, *rest]: + pass + case {"k": v, **kw}: + pass + lambda q=None: f"{q!r:>{w}}" +""") + + +def pre_order(node: ast.AST): + yield node + for child in ast.iter_child_nodes(node): + yield from pre_order(child) + + +@pytest.mark.parametrize( + "kinds", + [(ast.Name,), (ast.Call, ast.Attribute), (ast.Module,), (ast.Load,), ()], +) +def test_walk_of_types_is_the_filtered_pre_order(kinds): + tree = ast.parse(SOURCE) + assert [id(n) for n in walk_of_types(tree, kinds)] == [ + id(n) for n in pre_order(tree) if type(n) in kinds + ] + + +def test_walk_of_types_reads_fields_by_name(): + call = ast.Call(lineno=1, func=ast.Name(id="f", ctx=ast.Load()), args=[]) + assert [n.id for n in walk_of_types(call, (ast.Name,))] == ["f"] + + +SCOPES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda, ast.ClassDef) + + +def pruned_pre_order(node: ast.AST, prune: tuple[type, ...], root: bool = True): + yield node + if root or type(node) not in prune: + for child in ast.iter_child_nodes(node): + yield from pruned_pre_order(child, prune, root=False) + + +@pytest.mark.parametrize("kinds", [(ast.Name,), (ast.Return, *SCOPES), ()]) +def test_walk_of_types_stops_below_a_pruned_node(kinds): + tree = ast.parse(SOURCE) + for root in [tree, *walk_of_types(tree, SCOPES)]: + assert [id(n) for n in walk_of_types(root, kinds, SCOPES)] == [ + id(n) for n in pruned_pre_order(root, SCOPES) if type(n) in kinds + ] + + +def test_set_parents_matches_a_pre_order_loop(): + ours, theirs = ast.parse(SOURCE), ast.parse(SOURCE) + set_parents(ours) + for node in pre_order(theirs): + for child in ast.iter_child_nodes(node): + child.parent = node + assert not hasattr(ours, "parent") + position = {id(n): i for i, n in enumerate(pre_order(theirs))} + position |= {id(n): i for i, n in enumerate(pre_order(ours))} + assert [ + position[id(n.parent)] for n in pre_order(ours) if hasattr(n, "parent") + ] == [position[id(n.parent)] for n in pre_order(theirs) if hasattr(n, "parent")] + + +def test_set_parents_gives_a_shared_node_its_last_parent(): + shared = ast.Name(id="x", ctx=ast.Load()) + tree = ast.parse("f(a)\ng(b)") + tree.body[0].value.args[0] = shared + tree.body[1].value.args[0] = shared + set_parents(tree) + assert shared.parent is tree.body[1].value + + +def reference_events(node, kinds, bracket): + events = [] + + def descend(current): + if type(current) in bracket: + events.append((1, current)) + for child in ast.iter_child_nodes(current): + if type(child) in kinds: + events.append((0, child)) + else: + descend(child) + if type(current) in bracket: + events.append((2, current)) + + descend(node) + return events + + +EXPRESSIONS = tuple( + kind for kind in vars(ast).values() + if isinstance(kind, type) and issubclass(kind, ast.expr) and kind is not ast.expr +) + + +@pytest.mark.parametrize( + ("kinds", "bracket"), + [ + ((ast.Name, ast.Assign), EXPRESSIONS), + ((ast.Call,), (ast.Module, ast.Attribute)), + ((), EXPRESSIONS), + ((ast.FunctionDef,), ()), + ], +) +def test_walk_frontier_events_matches_a_recursive_descent(kinds, bracket): + tree = ast.parse(SOURCE) + for root in [tree, *walk_of_types(tree, (ast.Call, ast.Assign))]: + expected = reference_events(root, kinds, bracket) + assert [(e, id(n)) for e, n in walk_frontier_events(root, kinds, bracket)] == [ + (e, id(n)) for e, n in expected + ] + + +def test_walk_frontier_events_on_a_deep_chain(): + chain: ast.expr = ast.Name(id="a", ctx=ast.Load()) + for _ in range(50_000): + chain = ast.Attribute(value=chain, attr="b", ctx=ast.Load()) + events = walk_frontier_events(chain, (ast.Name,), EXPRESSIONS) + assert len(events) == 2 * 50_000 + 1 + + +def stripped(source: str) -> ast.Module: + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, (ast.expr, ast.stmt)) and len(node._fields) % 2: + for name in ("lineno", "col_offset", "end_lineno", "end_col_offset"): + if hasattr(node, name): + delattr(node, name) + return tree + + +def locations(tree: ast.AST): + return [ + tuple(getattr(n, name, "-") for name in n._attributes) for n in pre_order(tree) + ] + + +def test_fix_missing_locations_matches_the_stdlib(): + ours, theirs = stripped(SOURCE), stripped(SOURCE) + assert fix_missing_locations(ours) is ours + ast.fix_missing_locations(theirs) + assert locations(ours) == locations(theirs) + compile(ours, "", "exec") + + +def built() -> ast.Module: + return ast.Module( + body=[ + ast.Expr( + value=ast.Call( + lineno=3, + func=ast.Name(id="f", ctx=ast.Load()), + args=[], + keywords=[], + ) + ) + ], + type_ignores=[], + ) + + +def test_fix_missing_locations_fills_a_built_tree(): + ours, theirs = built(), built() + fix_missing_locations(ours) + ast.fix_missing_locations(theirs) + assert locations(ours) == locations(theirs) + assert ours.body[0].value.func.lineno == 3 + + +def test_fix_missing_locations_on_a_deep_chain(): + tree = ast.parse("x" + ".y" * 20_000) + for node in ast.walk(tree): + if isinstance(node, ast.Attribute): + del node.end_lineno + fix_missing_locations(tree) + assert tree.body[0].value.end_lineno == 1 + + +def pairs(left: ast.AST, right: ast.AST): + return [(a, b) for a in pre_order(left) for b in pre_order(right)] + + +def test_subtree_hashes_agree_with_ast_dump(): + old = ast.parse(SOURCE) + new = copy.deepcopy(old) + new.body[2].body[0].value.elt.func.id = "changed" + old_hashes, new_hashes = subtree_hashes(old), subtree_hashes(new) + for a, b in pairs(old, new): + assert (old_hashes[id(a)] == new_hashes[id(b)]) == ( + ast.dump(a) == ast.dump(b) + ), ast.dump(a) + + +@pytest.mark.parametrize( + ("left", "right"), [("x = 1", "x = True"), ("x = 1", "x = 1.0"), ("x = 0", "x = False")] +) +def test_subtree_hashes_tell_equal_hashing_constants_apart(left, right): + a, b = ast.parse(left), ast.parse(right) + assert subtree_hashes(a)[id(a)] != subtree_hashes(b)[id(b)] + + +@pytest.mark.parametrize( + ("built", "source"), + [ + (lambda: ast.Constant(value="x"), '"x"'), + (lambda: ast.Constant(value="x", kind=None), '"x"'), + (lambda: ast.Call(func=ast.Name(id="f"), args=[], keywords=[]), "f()"), + ], +) +def test_a_field_a_constructor_left_out_hashes_as_its_default(built, source): + node, parsed = built(), ast.parse(source).body[0].value + assert ast.dump(node) == ast.dump(parsed) + assert subtree_hash(node) == subtree_hash(parsed) + + +def test_a_field_a_constructor_left_out_is_not_a_child(): + node = ast.Constant(value="x") + assert walk_of_types(node, (ast.Constant,)) == [node] + assert walk_frontier_events(node, (), (ast.Constant,)) == [(1, node), (2, node)] + + +def test_subtree_hash_is_the_roots_entry(): + tree = ast.parse(SOURCE) + assert subtree_hash(tree) == subtree_hashes(tree)[id(tree)] + + +def test_subtree_hashes_ignore_dict_order(): + built = ast.Call(lineno=1, keywords=[], func=ast.Name(id="f", ctx=ast.Load()), args=[]) + parsed = ast.parse("f()").body[0].value + assert subtree_hashes(built)[id(built)] == subtree_hashes(parsed)[id(parsed)] + + +def test_refcount_neutral(): + tree = ast.parse(SOURCE) + set_parents(tree) + nodes = list(ast.walk(tree)) + gc.collect() + before = [sys.getrefcount(n) for n in nodes] + for _ in range(50): + walk_of_types(tree, (ast.Name, ast.Call)) + subtree_hashes(tree) + subtree_hash(tree) + set_parents(tree) + fix_missing_locations(tree) + walk_of_types(tree, (ast.Name,), SCOPES) + walk_frontier_events(tree, (ast.Name,), EXPRESSIONS) + gc.collect() + assert [sys.getrefcount(n) for n in nodes] == before + + +def name(id_: str = "x") -> ast.Name: + return ast.Name(id=id_, ctx=ast.Load()) + + +class Computed(ast.expr): + _fields = ("child",) + + @property + def child(self): + return name("computed") + + +def test_a_field_a_property_computes_survives_the_walk(): + node = Computed() + found = walk_of_types(node, (ast.Name,)) + assert [n.id for n in found] == ["computed"] + assert [n.id for n in walk_frontier_events(node, (ast.Name,), ())[0][1:]] == ["computed"] + assert len({subtree_hash(Computed()) for _ in range(3)}) == 1 + + +class Dropping(ast.expr): + _fields = ("first", "second") + + @property + def second(self): + self.__dict__.pop("first", None) + + +def test_a_field_read_may_drop_the_only_reference_to_a_child(): + node = Dropping() + node.first = name("first") + assert [n.id for n in walk_of_types(node, (ast.Name,))] == ["first"] + node = Dropping() + node.first = name("first") + set_parents(node) + + +def test_a_type_whose_address_is_reused_is_read_by_its_own_fields(): + for index in range(300): + kind = type(f"Kind{index}", (ast.expr,), {"_fields": (f"field{index}",)}) + node = kind() + setattr(node, f"field{index}", name(str(index))) + assert [n.id for n in walk_of_types(node, (ast.Name,))] == [str(index)] + del kind, node + gc.collect() + + +class MyAssign(ast.Assign): + pass + + +class Deeper(MyAssign): + pass + + +def test_a_subclass_of_any_depth_is_a_node(): + statement = Deeper(targets=[name("a")], value=name("b")) + tree = ast.Module(body=[statement], type_ignores=[]) + assert walk_of_types(tree, (ast.Name,)) == list(statement.targets) + [statement.value] + assert walk_of_types(tree, (Deeper,)) == [statement] + set_parents(tree) + assert statement.parent is tree + + +@pytest.mark.parametrize( + ("left", "right"), + [ + (lambda: float("nan"), lambda: float("nan")), + (lambda: (1, float("nan")), lambda: (1, float("nan"))), + (lambda: complex(float("nan"), 0), lambda: complex(float("nan"), 0)), + (lambda: 0.0, lambda: -0.0), + (lambda: 0j, lambda: -0j), + (lambda: [1], lambda: [1]), + (lambda: [1], lambda: [2]), + ], +) +def test_constant_values_hash_as_ast_dump_compares_them(left, right): + a, b = ast.Constant(value=left()), ast.Constant(value=right()) + assert (subtree_hash(a) == subtree_hash(b)) == (ast.dump(a) == ast.dump(b)) + + +class Defaulted(ast.AST): + _fields = ("child",) + child = name("default") + + +def test_a_node_with_no_instance_dict_reads_its_fields_by_attribute(): + node = Defaulted.__new__(Defaulted) + assert walk_of_types(node, (ast.Name,)) == list(ast.iter_child_nodes(node)) diff --git a/tests/test_frontier.py b/tests/test_frontier.py new file mode 100644 index 0000000..49e2bdf --- /dev/null +++ b/tests/test_frontier.py @@ -0,0 +1,119 @@ +"""`walk_frontier` and `walk_frontier_edges` must match a pure-Python reference: +the nodes `ast.NodeVisitor.generic_visit` dispatches to `visit_`, given +every other node falls through to the generic descent.""" + +from __future__ import annotations + +import ast +import gc +import sys +import textwrap + +import pytest + +from fast_walk import walk_frontier, walk_frontier_edges + +SOURCE = textwrap.dedent(""" + import x as y + def f(a, b=g(h.i)): + c = [k(l) for l in m.n.o] + if c: + return p.q(r, *s, t=u.v) + class W(X.Y): + z: Z = w.call() + lambda: q.attr +""") + +KINDS = [ + (ast.Call,), + (ast.Attribute,), + (ast.Call, ast.Attribute), + (ast.FunctionDef, ast.Name), + (ast.Load,), + (), +] + + +def reference(node: ast.AST, kinds: tuple[type, ...]): + found = [] + + def descend(parent: ast.AST): + for field, value in ast.iter_fields(parent): + items = value if isinstance(value, list) else [value] + for index, child in enumerate(items): + if not isinstance(child, ast.AST): + continue + if type(child) in kinds: + found.append(( + child, + parent, + field, + index if isinstance(value, list) else -1, + )) + else: + descend(child) + + descend(node) + return found + + +def identities(edges): + return [(id(n), id(p), f, i) for n, p, f, i in edges] + + +@pytest.mark.parametrize("kinds", KINDS) +def test_edges_match_the_generic_visit_order(kinds): + tree = ast.parse(SOURCE) + assert identities(walk_frontier_edges(tree, kinds)) == identities( + reference(tree, kinds) + ) + + +@pytest.mark.parametrize("kinds", KINDS) +def test_frontier_is_the_edge_nodes(kinds): + tree = ast.parse(SOURCE) + assert [id(n) for n in walk_frontier(tree, kinds)] == [ + id(edge[0]) for edge in walk_frontier_edges(tree, kinds) + ] + + +def test_the_root_is_never_included(): + tree = ast.parse("f(g())") + call = tree.body[0].value + assert [ast.unparse(n) for n in walk_frontier(call, (ast.Call,))] == ["g()"] + + +def test_a_match_is_not_descended(): + tree = ast.parse("a.b.c.d") + assert [ast.unparse(n) for n in walk_frontier(tree, (ast.Attribute,))] == [ + "a.b.c.d" + ] + + +def test_a_deep_chain_does_not_recurse(): + tree = ast.parse("x" + ".y" * 20_000) + assert [type(n) for n in walk_frontier(tree, (ast.Name,))] == [ast.Name] + + +def test_kinds_must_be_types(): + with pytest.raises(TypeError): + walk_frontier(ast.parse("x"), ("Name",)) + + +def test_refcount_neutral(): + tree = ast.parse(SOURCE) + nodes = list(ast.walk(tree)) + gc.collect() + before = [sys.getrefcount(n) for n in nodes] + for _ in range(50): + walk_frontier(tree, (ast.Call, ast.Attribute)) + walk_frontier_edges(tree, (ast.Call, ast.Attribute, ast.Name)) + gc.collect() + assert [sys.getrefcount(n) for n in nodes] == before + + +def test_fields_are_read_by_name(): + call = ast.Call(lineno=1, func=ast.Name(id="f", ctx=ast.Load()), args=[]) + assert [(type(n).__name__, f, i) for n, _, f, i in walk_frontier_edges(call, (ast.Name,))] == [ + ("Name", "func", -1) + ]