0.1.5

Author: adhami3310

.gitignore

@@ -41,6 +41,7 @@ htmlcov/
 .cache
 nosetests.xml
 coverage.xml
+.codspeed
 
 # Translations
 *.mo

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "fast-walk"
-version = "0.1.4"
+version = "0.1.5"
 edition = "2024"
 license = "MIT"
 readme = "README.md"
@@ -13,6 +13,8 @@ crate-type = ["cdylib"]
 
 [profile.release]
 debug = false
+lto = true        # Link-time optimization.
+codegen-units = 1 # Slower compilation but faster code. 
 
 [dependencies]
 pyo3 = "0.27.1"

Cargo.toml

@@ -12,7 +12,7 @@ pip install fast-walk
 
 ```python
 from fast_walk import walk
-import parse
+import ast
 
 code = """
 def hello(  x,y,   z  ):

README.md

@@ -4,12 +4,19 @@ build-backend = "maturin"
 
 [project]
 name = "fast-walk"
+description = "A fast reimplementation of ast.walk()"
 requires-python = ">=3.13"
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Programming Language :: Python :: 3 :: Only",
 ]
+license = { text = "MIT" }
+readme = "README.md"
 dynamic = ["version"]
 
 [tool.maturin]
@@ -23,3 +30,6 @@ cache-keys = [
     { file = "fast_walk.pyi" },
     { file = "src/**/*.rs" },
 ]
+
+[dependency-groups]
+dev = ["pytest>=8.4.2", "pytest-codspeed>=4.2.0"]

pyproject.toml

@@ -1,6 +1,24 @@
 use pyo3::types::{PyList, PyModule, PyTuple};
 use pyo3::{intern, prelude::*};
 
+fn getattr<'py>(
+    obj: &Bound<'py, PyAny>,
+    attr_name: &Bound<'py, PyAny>,
+) -> Option<Bound<'py, PyAny>> {
+    let py = obj.py();
+
+    let mut resp_ptr: *mut pyo3::ffi::PyObject = std::ptr::null_mut();
+    let attr_ptr = unsafe {
+        pyo3::ffi::PyObject_GetOptionalAttr(obj.as_ptr(), attr_name.as_ptr(), &mut resp_ptr)
+    };
+
+    if attr_ptr == 1 {
+        Some(unsafe { Bound::from_owned_ptr(py, resp_ptr) })
+    } else {
+        None
+    }
+}
+
 fn walk_node<'py>(
     node: Bound<'py, PyAny>,
     field_names: Bound<'py, PyTuple>,
@@ -10,18 +28,18 @@ fn walk_node<'py>(
 
     // Recursively walk through child nodes
     for field in field_names {
-        if let Ok(Some(child)) = node.getattr_opt(unsafe { field.cast_unchecked() }) {
+        if let Some(child) = getattr(&node, &field) {
             if child.is_exact_instance_of::<PyList>() {
                 for item in unsafe { child.cast_unchecked::<PyList>() } {
-                    if let Ok(Some(subfields)) = item.getattr_opt(intern!(item.py(), "_fields")) {
+                    if let Some(subfields) = getattr(&item, intern!(item.py(), "_fields")) {
                         walk_node(
                             item,
                             unsafe { subfields.cast_into_unchecked::<PyTuple>() },
                             result_list,
                         )?;
                     }
                 }
-            } else if let Ok(Some(subfields)) = child.getattr_opt(intern!(child.py(), "_fields")) {
+            } else if let Some(subfields) = getattr(&child, intern!(child.py(), "_fields")) {
                 walk_node(
                     child,
                     unsafe { subfields.cast_into_unchecked::<PyTuple>() },

src/lib.rs

@@ -0,0 +1,47 @@
+from ast import AST, parse
+from ast import walk as ast_walk
+from collections.abc import Callable, Iterable
+from pathlib import Path
+
+from pytest_codspeed import BenchmarkFixture
+from fast_walk import walk as fast_walk
+import pytest
+
+
+def python_walk_helper(node: AST, fields: tuple[str, ...], nodes: list[AST]):
+    nodes.append(node)
+    for field in fields:
+        value = getattr(node, field, None)
+        if type(value) is list:
+            for item in value:
+                if subfields := getattr(item, "_fields", None):
+                    python_walk_helper(item, subfields, nodes)
+        elif subfields := getattr(value, "_fields", None):
+            python_walk_helper(value, subfields, nodes)  # pyright: ignore[reportArgumentType]
+
+
+def python_walk(node: AST) -> list[AST]:
+    nodes: list[AST] = []
+    python_walk_helper(node, node._fields, nodes)
+    return nodes
+
+
+@pytest.mark.parametrize(
+    "algorithm",
+    [
+        ast_walk,
+        fast_walk,
+        python_walk,
+    ],
+)
+def test_walk(benchmark: BenchmarkFixture, algorithm: Callable[[AST], Iterable[AST]]):
+    import difflib
+
+    source_code = Path(difflib.__file__).read_text()
+    node = parse(source_code)
+
+    def run():
+        for _ in algorithm(node):
+            pass
+
+    benchmark(run)

tests/benchmarks.py

@@ -0,0 +1,11 @@
+if __name__ == "__main__":
+    from ast import parse
+    from pathlib import Path
+    from fast_walk import walk
+    import difflib
+
+    source_code = Path(difflib.__file__).read_text()
+    node = parse(source_code)
+
+    for _ in range(10_000):
+        walk(node)