Spaces:
Sleeping
Sleeping
File size: 3,472 Bytes
036a2db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """
parser.py — AST-based symbol extraction from Python source files.
Extracts functions, methods (including async), with class-aware naming.
"""
import ast
import logging
import os
from typing import Dict, List, Optional
from .models import Symbol
from ._warn_once import warn_syntax_error_once, check_and_warn_encoding
logger = logging.getLogger(__name__)
class _FunctionCollector(ast.NodeVisitor):
"""AST visitor that collects all function/method definitions."""
def __init__(self):
self.class_stack: list = []
self.collected: list = []
def visit_ClassDef(self, node):
self.class_stack.append(node.name)
for child in node.body:
self.visit(child)
self.class_stack.pop()
def visit_FunctionDef(self, node):
self._collect(node)
self.generic_visit(node)
def visit_AsyncFunctionDef(self, node):
self._collect(node)
self.generic_visit(node)
def _collect(self, node):
if self.class_stack:
name = ".".join(self.class_stack) + "." + node.name
else:
name = node.name
self.collected.append((name, node))
def extract_symbols(
filename: str,
repo_path: str,
broken_files: "Optional[List[str]]" = None,
) -> Dict[str, Symbol]:
"""
Parse a single Python file, return dict of symbol_id -> Symbol.
Symbol IDs look like: "./relative/path.py:ClassName.method_name"
If parsing fails and `broken_files` is provided (a list), the file's
relative path is appended to it so callers can distinguish "file failed
to parse" from "file legitimately has no functions."
"""
with open(filename, "rb") as f:
raw = f.read()
check_and_warn_encoding(logger, filename, raw)
source = raw.decode("utf-8", errors="ignore")
relative_file = "./" + os.path.relpath(filename, repo_path)
try:
tree = ast.parse(source)
except SyntaxError as e:
warn_syntax_error_once(logger, filename, e)
if broken_files is not None:
broken_files.append(relative_file)
return {}
collector = _FunctionCollector()
collector.visit(tree)
symbols = {}
for name, node in collector.collected:
symbol_id = f"{relative_file}:{name}"
code = ast.get_source_segment(source, node)
if code is None:
continue
symbols[symbol_id] = Symbol(
id=symbol_id,
file=filename,
name=name,
code=code,
lineno=node.lineno,
)
return symbols
def extract_all_symbols(
repo_path: str,
broken_files: "Optional[List[str]]" = None,
) -> Dict[str, Symbol]:
"""
Extract symbols from all Python files in a repository.
If `broken_files` is provided (a list), relative paths of any files
that failed to parse (SyntaxError) are appended to it.
"""
from .scanner import find_python_files
from .cache import SymbolCache
repo_path = os.path.abspath(repo_path)
all_symbols: Dict[str, Symbol] = {}
db_path = os.path.join(repo_path, ".diffcontext_cache.db")
with SymbolCache(db_path) as cache:
for filepath in find_python_files(repo_path):
def _parse(path: str) -> Dict[str, Symbol]:
return extract_symbols(path, repo_path, broken_files=broken_files)
all_symbols.update(cache.get_or_parse(filepath, _parse))
return all_symbols |