diff --git a/graphify/extractors/resolution.py b/graphify/extractors/resolution.py index 983b81be2..d0341ce77 100644 --- a/graphify/extractors/resolution.py +++ b/graphify/extractors/resolution.py @@ -10,6 +10,7 @@ _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 @@ -2232,14 +2250,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. @@ -2304,12 +2318,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). @@ -2415,7 +2429,7 @@ def visit(node, current_nid: str | None) -> None: 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(): diff --git a/tests/test_python_parse_memoization.py b/tests/test_python_parse_memoization.py new file mode 100644 index 000000000..0215f9877 --- /dev/null +++ b/tests/test_python_parse_memoization.py @@ -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"