perf(resolution): parse each Python file once across both resolution passes

The symbol-resolution facts pass and the cross-file import pass each parsed
the whole .py corpus, back to back, in the main process — every file read
and tree-sitter-parsed twice for nothing. _parse_python_tree now memoizes on
(path, mtime, size), so the second pass reuses the first pass's tree while a
file changed between watch runs still re-parses. On a 364-file self-corpus a
sequential extract drops from ~27s to ~24s, graph byte-identical.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JJbLfztSxm2tH5cJwBbe9q
(cherry picked from commit dcbab53879)
This commit is contained in:
abhay-codes07
2026-09-11 23:37:06 +05:30
committed by safishamsi
parent a572b88777
commit f077217191
2 changed files with 96 additions and 16 deletions
+30 -16
View File
@@ -10,6 +10,7 @@ from graphify.extractors.base import ( # noqa: F401
_make_id,
_read_text,
)
import functools
import hashlib
import json
import os
@@ -1938,13 +1939,30 @@ def _collect_js_symbol_resolution_facts(paths: list[Path], facts: _SymbolResolut
class_nid = _make_id(stem, class_name)
_ts_walk_class_members(node, source, path, class_nid, facts)
@functools.lru_cache(maxsize=2048)
def _parse_python_tree_cached(path_str: str, _mtime_ns: int, _size: int):
import tree_sitter_python as tspython
from tree_sitter import Language, Parser
source = Path(path_str).read_bytes()
parser = Parser(Language(tspython.language()))
return source, parser.parse(source).root_node
def _parse_python_tree(path: Path):
"""Parse one Python file to ``(source, root_node)``, memoized (#perf).
The Python symbol-resolution facts pass and the cross-file import pass each
parse the entire ``.py`` corpus, back to back, from the main process after
the workers return — so every file was tree-sitter-parsed (and read from
disk) twice for no reason. Keying the memo on ``(path, mtime_ns, size)``
lets the second pass reuse the first pass's tree while still re-parsing a
file that changed between runs (watch mode). Both passes only read the
tree, so sharing one parse is behaviour-preserving. Returns ``None`` on any
error, exactly as before — callers already treat that as "skip this file".
"""
try:
import tree_sitter_python as tspython
from tree_sitter import Language, Parser
source = path.read_bytes()
parser = Parser(Language(tspython.language()))
return source, parser.parse(source).root_node
st = path.stat()
return _parse_python_tree_cached(str(path), st.st_mtime_ns, st.st_size)
except Exception:
return None
@@ -2243,14 +2261,10 @@ def _resolve_cross_file_imports(
BasicAuth --uses--> Request [INFERRED]
"""
try:
import tree_sitter_python as tspython
from tree_sitter import Language, Parser
import tree_sitter_python # noqa: F401 (availability check only)
except ImportError:
return []
language = Language(tspython.language())
parser = Parser(language)
# Pass 1: _file_stem(path) → {ClassName: node_id}
# Keyed by directory-qualified stem (e.g. "auth_models") to avoid collisions
# when multiple files share the same filename in different directories.
@@ -2315,12 +2329,12 @@ def _resolve_cross_file_imports(
if not name_to_nid:
continue
# Parse imports from this file
try:
source = path.read_bytes()
tree = parser.parse(source)
except Exception:
# Parse imports from this file (shared with the facts pass via the
# mtime-keyed memo, so each .py is parsed once across both passes).
parsed = _parse_python_tree(path)
if parsed is None:
continue
source, root_node = parsed
# local_name -> target node id (local_name honours `import X as Y`, so a
# reference to the alias in the body still attributes correctly).
@@ -2426,7 +2440,7 @@ def _resolve_cross_file_imports(
for child in node.children:
visit(child, current_nid)
visit(tree.root_node, None)
visit(root_node, None)
for name, tgt_nid in import_targets.items():
for src_nid, line in ref_sources.get(name, {}).items():
+66
View File
@@ -0,0 +1,66 @@
"""_parse_python_tree memoizes across the two Python resolution passes (#perf).
The symbol-resolution facts pass and the cross-file import pass each parse the
whole .py corpus, back to back, in the main process so every file was read
and tree-sitter-parsed twice. The parse is now memoized on (path, mtime, size):
the second pass reuses the first pass's tree, a changed file still re-parses.
"""
import time
import pytest
pytest.importorskip("tree_sitter_python")
from graphify.extractors.resolution import _parse_python_tree
try:
from graphify.extractors.resolution import _parse_python_tree_cached
except ImportError: # pre-fix tree
_parse_python_tree_cached = None
needs_cache = pytest.mark.skipif(
_parse_python_tree_cached is None, reason="parse memo not present"
)
def test_parses_and_returns_source_and_root(tmp_path):
f = tmp_path / "m.py"
f.write_text("def foo():\n return 1\n", encoding="utf-8")
parsed = _parse_python_tree(f)
assert parsed is not None
source, root = parsed
assert b"def foo" in source
assert root.type == "module"
def test_missing_file_returns_none(tmp_path):
assert _parse_python_tree(tmp_path / "nope.py") is None
@needs_cache
def test_reuses_parse_within_a_run(tmp_path):
f = tmp_path / "m.py"
f.write_text("x = 1\n", encoding="utf-8")
_parse_python_tree_cached.cache_clear()
r1 = _parse_python_tree(f)
for _ in range(9):
_parse_python_tree(f)
info = _parse_python_tree_cached.cache_info()
assert info.misses == 1 and info.hits == 9, info
# Same cached (source, root) object handed back.
assert _parse_python_tree(f)[1] is r1[1]
@needs_cache
def test_edit_reparses(tmp_path):
f = tmp_path / "m.py"
f.write_text("x = 1\n", encoding="utf-8")
_parse_python_tree_cached.cache_clear()
src1, _ = _parse_python_tree(f)
assert b"x = 1" in src1
# A rewrite that changes size AND mtime must re-parse, not replay.
time.sleep(0.01)
f.write_text("x = 22222\ny = 3\n", encoding="utf-8")
src2, _ = _parse_python_tree(f)
assert b"y = 3" in src2, "edited file was served from a stale parse"