File size: 2,580 Bytes
b4291bc | 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 | import ast
from pathlib import Path
from backend.database.schemas import Chunk
def _source_segment(lines: list[str], start: int, end: int) -> str:
return "".join(lines[start - 1 : end]).strip()
def chunk_python_file(path: Path, repo_root: Path, repo_id: str) -> list[Chunk]:
text = path.read_text(encoding="utf-8", errors="ignore")
lines = text.splitlines(keepends=True)
rel = path.relative_to(repo_root).as_posix()
chunks: list[Chunk] = []
try:
tree = ast.parse(text)
except SyntaxError:
return [
Chunk(
id=f"{repo_id}:{rel}:1",
repo_id=repo_id,
path=rel,
language="python",
start_line=1,
end_line=max(1, len(lines)),
content=text,
)
]
imports = [
node
for node in tree.body
if isinstance(node, (ast.Import, ast.ImportFrom))
]
if imports:
start = min(node.lineno for node in imports)
end = max(getattr(node, "end_lineno", node.lineno) for node in imports)
chunks.append(
Chunk(
id=f"{repo_id}:{rel}:imports",
repo_id=repo_id,
path=rel,
language="python",
symbol="imports",
kind="imports",
start_line=start,
end_line=end,
content=_source_segment(lines, start, end),
)
)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
start = node.lineno
end = getattr(node, "end_lineno", start)
kind = "class" if isinstance(node, ast.ClassDef) else "function"
chunks.append(
Chunk(
id=f"{repo_id}:{rel}:{start}:{node.name}",
repo_id=repo_id,
path=rel,
language="python",
symbol=node.name,
kind=kind,
start_line=start,
end_line=end,
content=_source_segment(lines, start, end),
)
)
if not chunks:
chunks.append(
Chunk(
id=f"{repo_id}:{rel}:1",
repo_id=repo_id,
path=rel,
language="python",
start_line=1,
end_line=max(1, len(lines)),
content=text,
)
)
return chunks
|