182 lines
7.9 KiB
Python
182 lines
7.9 KiB
Python
from __future__ import annotations
|
|
|
|
from .shared import *
|
|
|
|
class SymbolIndexParsePythonMixin:
|
|
def _parse_file(
|
|
self,
|
|
*,
|
|
file_path: str,
|
|
language: str,
|
|
content: str,
|
|
) -> tuple[list[ParsedSymbol], list[ParsedReference]]:
|
|
if language == "python":
|
|
return self._parse_python_file(file_path=file_path, content=content)
|
|
return self._parse_typescript_file(file_path=file_path, content=content)
|
|
|
|
def _parse_python_file(self, *, file_path: str, content: str) -> tuple[list[ParsedSymbol], list[ParsedReference]]:
|
|
text = str(content or "")
|
|
if not text.strip():
|
|
return ([], [])
|
|
try:
|
|
tree = ast.parse(text)
|
|
except Exception:
|
|
return ([], [])
|
|
symbols: list[ParsedSymbol] = []
|
|
refs: list[ParsedReference] = []
|
|
|
|
class Visitor(ast.NodeVisitor):
|
|
def __init__(self):
|
|
self.class_stack: list[str] = []
|
|
self.func_stack: list[str] = []
|
|
|
|
def _current_source(self) -> str:
|
|
if self.func_stack:
|
|
return self.func_stack[-1]
|
|
if self.class_stack:
|
|
return self.class_stack[-1]
|
|
return ""
|
|
|
|
def _qualify(self, name: str) -> str:
|
|
chain = [*self.class_stack]
|
|
if self.func_stack:
|
|
chain.extend(self.func_stack[-1].split(".")[len(self.class_stack):])
|
|
chain.append(name)
|
|
uniq: list[str] = []
|
|
for item in chain:
|
|
if item and (not uniq or uniq[-1] != item):
|
|
uniq.append(item)
|
|
return ".".join(uniq) if uniq else name
|
|
|
|
def visit_ClassDef(self, node: ast.ClassDef) -> Any:
|
|
qualified = ".".join([*self.class_stack, node.name]) if self.class_stack else node.name
|
|
symbols.append(
|
|
ParsedSymbol(
|
|
symbol_type="class",
|
|
symbol_name=node.name,
|
|
qualified_name=qualified,
|
|
start_line=int(getattr(node, "lineno", 0) or 0),
|
|
end_line=int(getattr(node, "end_lineno", getattr(node, "lineno", 0)) or 0),
|
|
signature=f"class {node.name}",
|
|
docstring=str(ast.get_docstring(node) or ""),
|
|
payload={"file_path": file_path},
|
|
)
|
|
)
|
|
self.class_stack.append(node.name)
|
|
self.generic_visit(node)
|
|
self.class_stack.pop()
|
|
|
|
def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
|
|
self._visit_function_node(node=node, is_async=False)
|
|
|
|
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
|
|
self._visit_function_node(node=node, is_async=True)
|
|
|
|
def _visit_function_node(self, *, node: Any, is_async: bool) -> None:
|
|
if self.class_stack:
|
|
qualified = ".".join([*self.class_stack, node.name])
|
|
symbol_type = "method"
|
|
else:
|
|
qualified = node.name
|
|
symbol_type = "function"
|
|
arg_names = []
|
|
for arg in list(node.args.posonlyargs) + list(node.args.args):
|
|
arg_names.append(str(arg.arg))
|
|
if node.args.vararg is not None:
|
|
arg_names.append(f"*{node.args.vararg.arg}")
|
|
for arg in node.args.kwonlyargs:
|
|
arg_names.append(str(arg.arg))
|
|
if node.args.kwarg is not None:
|
|
arg_names.append(f"**{node.args.kwarg.arg}")
|
|
prefix = "async def" if is_async else "def"
|
|
signature = f"{prefix} {node.name}({', '.join(arg_names)})"
|
|
symbols.append(
|
|
ParsedSymbol(
|
|
symbol_type=symbol_type,
|
|
symbol_name=node.name,
|
|
qualified_name=qualified,
|
|
start_line=int(getattr(node, "lineno", 0) or 0),
|
|
end_line=int(getattr(node, "end_lineno", getattr(node, "lineno", 0)) or 0),
|
|
signature=signature[:1000],
|
|
docstring=str(ast.get_docstring(node) or ""),
|
|
payload={"file_path": file_path},
|
|
)
|
|
)
|
|
self.func_stack.append(qualified)
|
|
self.generic_visit(node)
|
|
self.func_stack.pop()
|
|
|
|
def visit_Import(self, node: ast.Import) -> Any:
|
|
source = self._current_source()
|
|
for alias in node.names:
|
|
target = str(alias.name or "").strip()
|
|
if not target:
|
|
continue
|
|
refs.append(
|
|
ParsedReference(
|
|
relation_type="imports",
|
|
target_symbol_name=target.split(".")[-1][:255],
|
|
target_qualified_name=target[:700],
|
|
source_qualified_name=source[:700],
|
|
start_line=int(getattr(node, "lineno", 0) or 0),
|
|
end_line=int(getattr(node, "end_lineno", getattr(node, "lineno", 0)) or 0),
|
|
payload={"alias": str(alias.asname or "")[:255] if alias.asname else None},
|
|
)
|
|
)
|
|
self.generic_visit(node)
|
|
|
|
def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:
|
|
source = self._current_source()
|
|
module = str(node.module or "").strip()
|
|
for alias in node.names:
|
|
name = str(alias.name or "").strip()
|
|
if not name:
|
|
continue
|
|
target_qn = f"{module}.{name}" if module else name
|
|
refs.append(
|
|
ParsedReference(
|
|
relation_type="imports",
|
|
target_symbol_name=name[:255],
|
|
target_qualified_name=target_qn[:700],
|
|
source_qualified_name=source[:700],
|
|
start_line=int(getattr(node, "lineno", 0) or 0),
|
|
end_line=int(getattr(node, "end_lineno", getattr(node, "lineno", 0)) or 0),
|
|
payload={"module": module[:255] if module else None},
|
|
)
|
|
)
|
|
self.generic_visit(node)
|
|
|
|
def visit_Call(self, node: ast.Call) -> Any:
|
|
source = self._current_source()
|
|
target = ""
|
|
if isinstance(node.func, ast.Name):
|
|
target = str(node.func.id or "").strip()
|
|
elif isinstance(node.func, ast.Attribute):
|
|
target = str(node.func.attr or "").strip()
|
|
if target:
|
|
refs.append(
|
|
ParsedReference(
|
|
relation_type="calls",
|
|
target_symbol_name=target[:255],
|
|
target_qualified_name=target[:700],
|
|
source_qualified_name=source[:700],
|
|
start_line=int(getattr(node, "lineno", 0) or 0),
|
|
end_line=int(getattr(node, "end_lineno", getattr(node, "lineno", 0)) or 0),
|
|
payload={},
|
|
)
|
|
)
|
|
self.generic_visit(node)
|
|
|
|
Visitor().visit(tree)
|
|
return (symbols[:5000], refs[:12000])
|
|
|
|
@staticmethod
|
|
def _detect_language(path: Path) -> str:
|
|
ext = path.suffix.lower()
|
|
if ext == ".py":
|
|
return "python"
|
|
if ext in {".ts", ".tsx", ".js", ".jsx"}:
|
|
return "typescript"
|
|
return ""
|
|
|