backend/chunking/python_chunker.py

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