perf: cache call graph call nodes (#1215)

Author: mldangelo-oai

packages/modelaudit-picklescan/src/modelaudit_picklescan/call_graph.py

@@ -757,6 +757,7 @@ def _clear_source_sensitive_caches_now() -> None:
         _source_function_context,
         _source_class_context,
         _constructor_parameter_self_attribute_targets,
+        _iter_call_nodes,
         _collect_function_import_aliases,
         _parameter_controlled_names,
         _can_invoke_function_with_positional_args,
@@ -2602,6 +2603,7 @@ def _function_instance_alias_value(
     return _resolve_class_target(call_target)
 
 
+@lru_cache(maxsize=4096)
 def _iter_call_nodes(function_node: ast.FunctionDef | ast.AsyncFunctionDef) -> tuple[ast.Call, ...]:
     calls: list[ast.Call] = []
 

packages/modelaudit-picklescan/tests/test_call_graph_import_statements.py

@@ -14,6 +14,7 @@
 from importlib.machinery import ModuleSpec
 from importlib.util import find_spec
 from pathlib import Path
+from typing import Any
 
 import pytest
 
@@ -79,6 +80,7 @@ def _clear_call_graph_caches() -> None:
         "_resolve_function_target",
         "_resolve_module_source",
         "_safe_call_graph_entrypoints",
+        "_iter_call_nodes",
         "_collect_function_import_aliases",
         "_parameter_controlled_names",
         "_split_function_name",
@@ -140,6 +142,37 @@ def _env_without_pythonpath() -> dict[str, str]:
     return {key: value for key, value in os.environ.items() if key != "PYTHONPATH"}
 
 
+def test_iter_call_nodes_reuses_cached_walk(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = ast.parse(
+        """
+def bridge(target, command):
+    callback = getattr(target, command)
+    callback(command)
+"""
+    )
+    bridge = module.body[0]
+    assert isinstance(bridge, ast.FunctionDef)
+
+    visits = 0
+    original_visit = ast.NodeVisitor.visit
+
+    def counting_visit(self: ast.NodeVisitor, node: ast.AST) -> Any:
+        nonlocal visits
+        visits += 1
+        return original_visit(self, node)
+
+    monkeypatch.setattr(ast.NodeVisitor, "visit", counting_visit)
+    call_graph._iter_call_nodes.cache_clear()
+
+    first = call_graph._iter_call_nodes(bridge)
+    first_visit_count = visits
+    second = call_graph._iter_call_nodes(bridge)
+
+    assert first == second
+    assert first_visit_count > 0
+    assert visits == first_visit_count
+
+
 def test_collect_function_import_aliases_reuses_cached_walk(monkeypatch: pytest.MonkeyPatch) -> None:
     module = ast.parse(
         """