165 lines
6.7 KiB
Python
165 lines
6.7 KiB
Python
"""Граф зависимостей кода: calls/imports/inherits через tree-sitter.
|
|
|
|
Строит ориентированный граф зависимостей между символами кода:
|
|
- calls: функция A вызывает функцию B
|
|
- imports: файл импортирует модуль (src = имя файла без расширения)
|
|
- inherits: класс B наследует класс A
|
|
|
|
Извлекает из tree-sitter AST и сохраняет в pm_code_dependency (PostgreSQL).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from .code_symbols import _parser
|
|
|
|
_FUNCTION_TYPES = ("function_definition", "function_item", "function_declaration", "method_definition")
|
|
_CLASS_TYPES = ("class_definition", "class_declaration")
|
|
|
|
|
|
def _sym_name(node, sb: bytes) -> str | None:
|
|
n = node.child_by_field_name("name")
|
|
return sb[n.start_byte:n.end_byte].decode("utf-8", "replace") if n else None
|
|
|
|
|
|
def _call_target(node, sb: bytes) -> str | None:
|
|
"""Имя вызываемой функции из call-узла."""
|
|
fn = node.child_by_field_name("function")
|
|
if fn is None:
|
|
return None
|
|
if fn.type == "identifier":
|
|
return sb[fn.start_byte:fn.end_byte].decode("utf-8", "replace")
|
|
if fn.type == "attribute":
|
|
parts = [sb[c.start_byte:c.end_byte].decode("utf-8", "replace")
|
|
for c in fn.children if c.type == "identifier"]
|
|
return parts[-1] if parts else None
|
|
return None
|
|
|
|
|
|
def _collect_calls(node, owner, sb, add_fn):
|
|
"""Рекурсивно собрать вызовы внутри тела функции/метода."""
|
|
if node.type == "call":
|
|
t = _call_target(node, sb)
|
|
if t:
|
|
add_fn(owner, t, "calls")
|
|
for ch in node.children:
|
|
_collect_calls(ch, owner, sb, add_fn)
|
|
|
|
|
|
def _collect_inherits(node, owner, sb, add_fn):
|
|
"""Найти базовые классы (superclasses/heritage_clause)."""
|
|
for ch in node.children:
|
|
if ch.type in ("argument_list", "superclasses", "heritage_clause", "base_list"):
|
|
for sub in _iter(ch):
|
|
if sub.type == "identifier":
|
|
add_fn(owner, sb[sub.start_byte:sub.end_byte].decode("utf-8", "replace"), "inherits")
|
|
elif sub.type == "attribute":
|
|
parts = [sb[p.start_byte:p.end_byte].decode("utf-8", "replace")
|
|
for p in sub.children if p.type == "identifier"]
|
|
if parts:
|
|
add_fn(owner, parts[-1], "inherits")
|
|
|
|
|
|
def _collect_imports(root, module_name, sb, add_fn):
|
|
"""Собрать импорты (import X / from Y import Z)."""
|
|
for node in _iter(root):
|
|
if node.type == "import_statement":
|
|
for ch in node.children:
|
|
if ch.type in ("dotted_name", "identifier"):
|
|
add_fn(module_name, sb[ch.start_byte:ch.end_byte].decode("utf-8", "replace"), "imports")
|
|
break
|
|
elif node.type == "import_from_statement":
|
|
mod = node.child_by_field_name("module_name")
|
|
if mod is not None:
|
|
target = sb[mod.start_byte:mod.end_byte].decode("utf-8", "replace")
|
|
# отфильтровать артефакты: "." (relative import), "..", пустые
|
|
if target not in (".", "..", ""):
|
|
add_fn(module_name, target, "imports")
|
|
|
|
|
|
def _iter(node):
|
|
yield node
|
|
for ch in node.children:
|
|
yield from _iter(ch)
|
|
|
|
|
|
def extract_dependencies(source: str, language: str, repo: str = "", path: str = "") -> list[dict]:
|
|
"""Вернуть зависимости [{repo,path,language,src,dst,kind}] из исходника."""
|
|
parser, _ = _parser(language)
|
|
sb = source.encode("utf-8")
|
|
tree = parser.parse(sb)
|
|
deps: list[dict] = []
|
|
seen: set[tuple] = set()
|
|
|
|
def add(src: str, dst: str, kind: str):
|
|
if not dst:
|
|
return
|
|
key = (src, dst, kind)
|
|
if key not in seen:
|
|
seen.add(key)
|
|
deps.append({"repo": repo, "path": path, "language": language,
|
|
"src": src, "dst": dst, "kind": kind})
|
|
|
|
module_name = path.rsplit("/", 1)[-1].rsplit(".", 1)[0] if path else "__module__"
|
|
|
|
def walk(node):
|
|
t = node.type
|
|
if t in _FUNCTION_TYPES:
|
|
owner = _sym_name(node, sb)
|
|
body = node.child_by_field_name("body")
|
|
if body is not None and owner:
|
|
_collect_calls(body, owner, sb, add)
|
|
return # не спускаемся в вложенные scope как отдельные
|
|
if t in _CLASS_TYPES:
|
|
owner = _sym_name(node, sb)
|
|
if owner:
|
|
_collect_inherits(node, owner, sb, add)
|
|
body = node.child_by_field_name("body")
|
|
if body is not None:
|
|
for ch in body.children:
|
|
walk(ch) # методы класса — вложенные function_definition
|
|
return
|
|
for ch in node.children:
|
|
walk(ch)
|
|
|
|
_collect_imports(tree.root_node, module_name, sb, add)
|
|
walk(tree.root_node)
|
|
return deps
|
|
|
|
|
|
def store_dependencies(conn, source: str, language: str, repo: str = "", path: str = "",
|
|
sync_memgraph: bool = True) -> int:
|
|
"""Извлечь и сохранить зависимости кода (+ рёбра Memgraph).
|
|
|
|
Возвращает число сохранённых.
|
|
"""
|
|
from . import memgraph_store
|
|
deps = extract_dependencies(source, language, repo, path)
|
|
n = 0
|
|
with conn.cursor() as cur:
|
|
for d in deps:
|
|
cur.execute(
|
|
"""INSERT INTO pm_code_dependency (repo, path, language, src, dst, kind)
|
|
VALUES (%s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (repo, path, src, dst, kind) DO NOTHING""",
|
|
(d["repo"], d["path"], d["language"], d["src"], d["dst"], d["kind"]),
|
|
)
|
|
n += 1
|
|
if sync_memgraph:
|
|
try:
|
|
memgraph_store.upsert_code_dependency(
|
|
d["src"], d["dst"], d["kind"], repo=d["repo"], path=d["path"],
|
|
)
|
|
except Exception:
|
|
pass
|
|
conn.commit()
|
|
return n
|
|
|
|
|
|
def store_file_dependencies(conn, path: str, repo: str = "") -> int:
|
|
"""По расширению файла определить язык и сохранить зависимости."""
|
|
ext = path.rsplit(".", 1)[-1].lower() if "." in path else ""
|
|
lang = {".py": "python", ".js": "javascript", ".ts": "typescript", ".rs": "rust"}.get("." + ext)
|
|
if lang is None:
|
|
return 0
|
|
source = open(path, encoding="utf-8").read()
|
|
return store_dependencies(conn, source, lang, repo, path)
|