chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,303 @@
|
||||
"""Post-build Jedi enrichment for Python call resolution.
|
||||
|
||||
After tree-sitter parsing, many method calls on lowercase-receiver variables
|
||||
are dropped (e.g. ``svc.authenticate()`` where ``svc = factory()``). Jedi
|
||||
can resolve these by tracing return types across files.
|
||||
|
||||
This module runs as a post-build step: it re-walks Python ASTs to find
|
||||
dropped calls, uses ``jedi.Script.goto()`` to resolve them, and adds the
|
||||
resulting CALLS edges to the graph database.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from .parser import CodeParser, EdgeInfo
|
||||
from .parser import _is_test_file as _parser_is_test_file
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SELF_NAMES = frozenset({"self", "cls", "super"})
|
||||
|
||||
|
||||
def enrich_jedi_calls(store, repo_root: Path) -> dict:
|
||||
"""Resolve untracked Python method calls via Jedi.
|
||||
|
||||
Walks Python files, finds ``receiver.method()`` calls that tree-sitter
|
||||
dropped (lowercase receiver, not self/cls), resolves them with Jedi,
|
||||
and inserts new CALLS edges.
|
||||
|
||||
Returns stats dict with ``resolved`` count.
|
||||
"""
|
||||
try:
|
||||
import jedi
|
||||
except ImportError:
|
||||
logger.info("Jedi not installed, skipping Python enrichment")
|
||||
return {"skipped": True, "reason": "jedi not installed"}
|
||||
|
||||
repo_root = Path(repo_root).resolve()
|
||||
|
||||
# Get Python files from the graph — skip early if none
|
||||
all_files = store.get_all_files()
|
||||
py_files = [f for f in all_files if f.endswith(".py")]
|
||||
|
||||
if not py_files:
|
||||
return {"resolved": 0, "files": 0}
|
||||
|
||||
# Scope the Jedi project to Python-only directories to avoid scanning
|
||||
# non-Python files (e.g. node_modules, TS sources). This matters for
|
||||
# polyglot monorepos where jedi.Project(path=repo_root) would scan
|
||||
# thousands of irrelevant files during initialization.
|
||||
py_dirs = sorted({str(Path(f).parent) for f in py_files})
|
||||
common_py_root = Path(os.path.commonpath(py_dirs)) if py_dirs else repo_root
|
||||
if not str(common_py_root).startswith(str(repo_root)):
|
||||
common_py_root = repo_root
|
||||
project = jedi.Project(
|
||||
path=str(common_py_root),
|
||||
added_sys_path=[str(repo_root)],
|
||||
smart_sys_path=False,
|
||||
)
|
||||
|
||||
# Pre-parse all Python files to find which ones have pending method calls.
|
||||
# This avoids expensive Jedi Script creation for files with nothing to resolve.
|
||||
parser = CodeParser()
|
||||
ts_parser = parser._get_parser("python")
|
||||
if not ts_parser:
|
||||
return {"resolved": 0, "files": 0}
|
||||
|
||||
# Build set of method names that actually exist in project code.
|
||||
# No point asking Jedi to resolve `logger.getLogger()` if no project
|
||||
# file defines a function called `getLogger`.
|
||||
project_func_names = {
|
||||
r["name"]
|
||||
for r in store._conn.execute(
|
||||
"SELECT DISTINCT name FROM nodes WHERE kind IN ('Function', 'Test')"
|
||||
).fetchall()
|
||||
}
|
||||
|
||||
files_with_pending: list[tuple[str, bytes, list]] = []
|
||||
total_skipped = 0
|
||||
for file_path in py_files:
|
||||
try:
|
||||
source = Path(file_path).read_bytes()
|
||||
except (OSError, PermissionError):
|
||||
continue
|
||||
tree = ts_parser.parse(source)
|
||||
is_test = _parser_is_test_file(file_path)
|
||||
pending = _find_untracked_method_calls(tree.root_node, is_test)
|
||||
if pending:
|
||||
# Only keep calls whose method name exists in project code
|
||||
filtered = [p for p in pending if p[2] in project_func_names]
|
||||
total_skipped += len(pending) - len(filtered)
|
||||
if filtered:
|
||||
files_with_pending.append((file_path, source, filtered))
|
||||
|
||||
if not files_with_pending:
|
||||
return {"resolved": 0, "files": 0}
|
||||
|
||||
logger.debug(
|
||||
"Jedi: %d/%d Python files have pending calls (%d calls skipped — no project target)",
|
||||
len(files_with_pending), len(py_files), total_skipped,
|
||||
)
|
||||
|
||||
resolved_count = 0
|
||||
files_enriched = 0
|
||||
errors = 0
|
||||
|
||||
for file_path, source, pending in files_with_pending:
|
||||
source_text = source.decode("utf-8", errors="replace")
|
||||
|
||||
# Get existing CALLS edges for this file to skip duplicates
|
||||
existing = set()
|
||||
for edge in _get_file_call_edges(store, file_path):
|
||||
existing.add((edge.source_qualified, edge.line))
|
||||
|
||||
# Get function nodes from DB for enclosing-function lookup
|
||||
func_nodes = [
|
||||
n for n in store.get_nodes_by_file(file_path)
|
||||
if n.kind in ("Function", "Test")
|
||||
]
|
||||
|
||||
# Create Jedi script once per file
|
||||
try:
|
||||
script = jedi.Script(source_text, path=file_path, project=project)
|
||||
except Exception as e:
|
||||
logger.debug("Jedi failed to load %s: %s", file_path, e)
|
||||
errors += 1
|
||||
continue
|
||||
|
||||
file_resolved = 0
|
||||
for jedi_line, col, _method_name, _enclosing_name in pending:
|
||||
# Find enclosing function qualified name
|
||||
enclosing = _find_enclosing(func_nodes, jedi_line)
|
||||
if not enclosing:
|
||||
enclosing = file_path # module-level
|
||||
|
||||
# Skip if we already have a CALLS edge from this source at this line
|
||||
if (enclosing, jedi_line) in existing:
|
||||
continue
|
||||
|
||||
# Ask Jedi to resolve
|
||||
try:
|
||||
names = script.goto(jedi_line, col)
|
||||
except Exception: # nosec B112 - Jedi may fail on malformed code
|
||||
continue
|
||||
|
||||
if not names:
|
||||
continue
|
||||
|
||||
name = names[0]
|
||||
if not name.module_path:
|
||||
continue
|
||||
|
||||
module_path = Path(name.module_path).resolve()
|
||||
|
||||
# Only emit edges for project-internal definitions
|
||||
try:
|
||||
module_path.relative_to(repo_root)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Build qualified target: file_path::Class.method or file_path::func
|
||||
target_file = str(module_path)
|
||||
parent = name.parent()
|
||||
if parent and parent.type == "class":
|
||||
target = f"{target_file}::{parent.name}.{name.name}"
|
||||
else:
|
||||
target = f"{target_file}::{name.name}"
|
||||
|
||||
store.upsert_edge(EdgeInfo(
|
||||
kind="CALLS",
|
||||
source=enclosing,
|
||||
target=target,
|
||||
file_path=file_path,
|
||||
line=jedi_line,
|
||||
))
|
||||
existing.add((enclosing, jedi_line))
|
||||
file_resolved += 1
|
||||
|
||||
if file_resolved:
|
||||
files_enriched += 1
|
||||
resolved_count += file_resolved
|
||||
|
||||
if resolved_count:
|
||||
store.commit()
|
||||
logger.info(
|
||||
"Jedi enrichment: resolved %d calls in %d files",
|
||||
resolved_count, files_enriched,
|
||||
)
|
||||
|
||||
return {
|
||||
"resolved": resolved_count,
|
||||
"files": files_enriched,
|
||||
"errors": errors,
|
||||
}
|
||||
|
||||
|
||||
def _get_file_call_edges(store, file_path: str):
|
||||
"""Get all CALLS edges originating from a file."""
|
||||
conn = store._conn
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM edges WHERE file_path = ? AND kind = 'CALLS'",
|
||||
(file_path,),
|
||||
).fetchall()
|
||||
from .graph import GraphEdge
|
||||
return [
|
||||
GraphEdge(
|
||||
id=r["id"], kind=r["kind"],
|
||||
source_qualified=r["source_qualified"],
|
||||
target_qualified=r["target_qualified"],
|
||||
file_path=r["file_path"], line=r["line"],
|
||||
extra={},
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
|
||||
|
||||
def _find_enclosing(func_nodes, line: int) -> Optional[str]:
|
||||
"""Find the qualified name of the function enclosing a given line."""
|
||||
best = None
|
||||
best_span = float("inf")
|
||||
for node in func_nodes:
|
||||
if node.line_start <= line <= node.line_end:
|
||||
span = node.line_end - node.line_start
|
||||
if span < best_span:
|
||||
best = node.qualified_name
|
||||
best_span = span
|
||||
return best
|
||||
|
||||
|
||||
def _find_untracked_method_calls(root, is_test_file: bool = False):
|
||||
"""Walk Python AST to find method calls the parser would have dropped.
|
||||
|
||||
Returns list of (jedi_line, col, method_name, enclosing_func_name) tuples.
|
||||
Jedi_line is 1-indexed, col is 0-indexed.
|
||||
"""
|
||||
results: list[tuple[int, int, str, Optional[str]]] = []
|
||||
_walk_calls(root, results, is_test_file, enclosing_func=None)
|
||||
return results
|
||||
|
||||
|
||||
def _walk_calls(node, results, is_test_file, enclosing_func):
|
||||
"""Recursively walk AST collecting dropped method calls."""
|
||||
# Track enclosing function scope
|
||||
if node.type == "function_definition":
|
||||
name = None
|
||||
for child in node.children:
|
||||
if child.type == "identifier":
|
||||
name = child.text.decode("utf-8", errors="replace")
|
||||
break
|
||||
for child in node.children:
|
||||
_walk_calls(child, results, is_test_file, name or enclosing_func)
|
||||
return
|
||||
|
||||
if node.type == "decorated_definition":
|
||||
for child in node.children:
|
||||
_walk_calls(child, results, is_test_file, enclosing_func)
|
||||
return
|
||||
|
||||
# Check for call expressions with attribute access
|
||||
if node.type == "call":
|
||||
first = node.children[0] if node.children else None
|
||||
if first and first.type == "attribute":
|
||||
_check_dropped_call(first, results, is_test_file, enclosing_func)
|
||||
|
||||
for child in node.children:
|
||||
_walk_calls(child, results, is_test_file, enclosing_func)
|
||||
|
||||
|
||||
def _check_dropped_call(attr_node, results, is_test_file, enclosing_func):
|
||||
"""Check if an attribute-based call was dropped by the parser."""
|
||||
children = attr_node.children
|
||||
if len(children) < 2:
|
||||
return
|
||||
|
||||
receiver = children[0]
|
||||
# Only handle simple identifier receivers
|
||||
if receiver.type != "identifier":
|
||||
return
|
||||
|
||||
receiver_text = receiver.text.decode("utf-8", errors="replace")
|
||||
|
||||
# The parser keeps: self/cls/super calls and uppercase-receiver calls
|
||||
# The parser keeps: calls handled by typed-var enrichment (but those are
|
||||
# separate edges -- we check for duplicates via existing-edge set)
|
||||
if receiver_text in _SELF_NAMES:
|
||||
return
|
||||
if receiver_text[:1].isupper():
|
||||
return
|
||||
if is_test_file:
|
||||
return # test files already track all calls
|
||||
|
||||
# Find the method name identifier
|
||||
method_node = children[-1]
|
||||
if method_node.type != "identifier":
|
||||
return
|
||||
|
||||
row, col = method_node.start_point # 0-indexed
|
||||
method_name = method_node.text.decode("utf-8", errors="replace")
|
||||
results.append((row + 1, col, method_name, enclosing_func))
|
||||
Reference in New Issue
Block a user