wehub-resource-sync
304 lines
10 KiB
Python
304 lines
10 KiB
Python
"""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))
|