109 lines
4.9 KiB
Python
109 lines
4.9 KiB
Python
"""Извлечение символов кода через tree-sitter (Фаза 3).
|
||
|
||
Грамматики ставятся отдельными пакетами: tree-sitter-python, tree-sitter-javascript,
|
||
tree-sitter-rust и т.д. Каждый пакет экспортирует language().
|
||
"""
|
||
from __future__ import annotations
|
||
import hashlib
|
||
|
||
# Язык -> (имя модуля, признак function/class узлов)
|
||
_LANGS = {
|
||
"python": ("tree_sitter_python", {"function_definition", "class_definition"}),
|
||
"javascript": ("tree_sitter_javascript", {"function_declaration", "class_declaration", "method_definition"}),
|
||
"typescript": ("tree_sitter_javascript", {"function_declaration", "class_declaration", "method_definition"}),
|
||
"rust": ("tree_sitter_rust", {"function_item", "struct_item", "impl_item", "trait_item", "enum_item"}),
|
||
}
|
||
|
||
|
||
def _parser(language: str):
|
||
"""Вернуть (Parser, kind_nodes) для языка."""
|
||
from tree_sitter import Language, Parser
|
||
spec = _LANGS.get(language)
|
||
if spec is None:
|
||
raise ValueError(f"unsupported language: {language}")
|
||
mod_name, kind_nodes = spec
|
||
mod = __import__(mod_name)
|
||
lang = Language(mod.language())
|
||
return Parser(lang), kind_nodes
|
||
|
||
|
||
def extract_symbols(source: str, language: str, repo: str = "", path: str = "") -> list[dict]:
|
||
"""Вернуть символы (function/class/method/type/const) из исходника.
|
||
|
||
tree-sitter индексирует байтами UTF-8, поэтому слайсим из bytes (src),
|
||
а не из str (source) — иначе кириллица в комментариях сбивает индексы.
|
||
"""
|
||
parser, kind_nodes = _parser(language)
|
||
src = source.encode("utf-8")
|
||
tree = parser.parse(src)
|
||
symbols: list[dict] = []
|
||
|
||
def walk(node):
|
||
if node.type in kind_nodes:
|
||
name_node = node.child_by_field_name("name")
|
||
name = src[name_node.start_byte:name_node.end_byte].decode("utf-8", "replace") if name_node else "?"
|
||
# сигнатура: от начала узла до первого тела/блока (первые 200 байт)
|
||
sig = src[node.start_byte:min(node.end_byte, node.start_byte + 200)].decode("utf-8", "replace")
|
||
if "function" in node.type or node.type in ("function_item", "function_declaration", "method_definition"):
|
||
kind = "function"
|
||
elif "class" in node.type or node.type in ("struct_item", "impl_item", "trait_item"):
|
||
kind = "class"
|
||
elif "enum" in node.type:
|
||
kind = "type"
|
||
else:
|
||
kind = "type"
|
||
symbols.append({
|
||
"repo": repo, "path": path, "language": language,
|
||
"symbol_name": name, "kind": kind,
|
||
"signature": sig,
|
||
"body_hash": hashlib.sha256(src[node.start_byte:node.end_byte]).hexdigest(),
|
||
"line_start": node.start_point[0] + 1, "line_end": node.end_point[0] + 1,
|
||
})
|
||
for child in node.children:
|
||
walk(child)
|
||
|
||
walk(tree.root_node)
|
||
return symbols
|
||
|
||
|
||
def store_symbols(conn, source: str, language: str, repo: str = "", path: str = "",
|
||
sync_memgraph: bool = True) -> int:
|
||
"""Извлечь и сохранить символы кода в pm_code_symbol (+ узлы Memgraph).
|
||
|
||
Возвращает число сохранённых.
|
||
"""
|
||
from . import memgraph_store
|
||
syms = extract_symbols(source, language, repo, path)
|
||
n = 0
|
||
with conn.cursor() as cur:
|
||
for s in syms:
|
||
cur.execute(
|
||
"""INSERT INTO pm_code_symbol (repo, path, language, symbol_name, kind,
|
||
signature, body_hash, line_start, line_end)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||
ON CONFLICT DO NOTHING""",
|
||
(s["repo"], s["path"], s["language"], s["symbol_name"], s["kind"],
|
||
s["signature"], s["body_hash"], s["line_start"], s["line_end"]),
|
||
)
|
||
n += 1
|
||
if sync_memgraph:
|
||
try:
|
||
memgraph_store.upsert_code_symbol(
|
||
s["symbol_name"], repo=s["repo"], path=s["path"],
|
||
language=s["language"], kind=s["kind"],
|
||
)
|
||
except Exception:
|
||
pass
|
||
conn.commit()
|
||
return n
|
||
|
||
|
||
def store_file(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_symbols(conn, source, lang, repo, path)
|