btl/annotate.py

import ast
from dataclasses import dataclass
from typing import Literal

from .model import ModelUnavailableError, generate_comment, generate_file_summary


ModelChoice = Literal["base", "tuned"]


@dataclass(frozen=True)
class BlockInfo:
    kind: str
    name: str
    lineno: int
    source: str


@dataclass(frozen=True)
class AnnotationResult:
    summary: str
    annotated_source: str
    status: str
    valid: bool
    block_count: int
    notes: tuple[str, ...] = ()


def parse_python(source: str) -> ast.Module:
    return ast.parse(source)


def strip_docstrings(node: ast.AST) -> ast.AST:
    copied = ast.fix_missing_locations(ast.parse(ast.unparse(node)))
    for child in ast.walk(copied):
        if isinstance(child, (ast.Module, ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            if (
                child.body
                and isinstance(child.body[0], ast.Expr)
                and isinstance(child.body[0].value, ast.Constant)
                and isinstance(child.body[0].value.value, str)
            ):
                child.body = child.body[1:]
    return copied


def semantic_ast_dump(source: str) -> str:
    tree = parse_python(source)
    stripped = strip_docstrings(tree)
    return ast.dump(stripped, include_attributes=False)


def collect_blocks(tree: ast.Module, source: str) -> list[BlockInfo]:
    blocks: list[BlockInfo] = []
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            block_source = ast.get_source_segment(source, node) or ast.unparse(node)
            blocks.append(BlockInfo("class", node.name, node.lineno, block_source))
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            kind = "async function" if isinstance(node, ast.AsyncFunctionDef) else "function"
            block_source = ast.get_source_segment(source, node) or ast.unparse(node)
            blocks.append(BlockInfo(kind, node.name, node.lineno, block_source))
    return sorted(blocks, key=lambda item: item.lineno)


def summarize_blocks(blocks: list[BlockInfo]) -> str:
    if not blocks:
        return "This Python file has no classes or functions to annotate yet."
    names = ", ".join(f"{block.kind} `{block.name}`" for block in blocks[:8])
    overflow = "" if len(blocks) <= 8 else f", plus {len(blocks) - 8} more"
    return f"This file defines {names}{overflow}."


def generate_summary(source: str, blocks: list[BlockInfo], model_choice: ModelChoice = "base") -> tuple[str, list[str]]:
    notes: list[str] = []
    try:
        summary = generate_file_summary(source, variant=model_choice)
        if summary:
            return summary, notes
        notes.append("file summary: model returned an empty summary.")
    except ModelUnavailableError:
        raise
    except Exception as exc:
        notes.append(f"file summary: model generation failed ({type(exc).__name__}).")
    return summarize_blocks(blocks), notes


def insert_comments(source: str, comments: dict[int, str]) -> str:
    lines = source.splitlines()
    inserts: dict[int, list[str]] = {}
    for lineno, comment_text in comments.items():
        indent = len(lines[lineno - 1]) - len(lines[lineno - 1].lstrip())
        inserts.setdefault(lineno - 1, []).append(" " * indent + comment_text)

    annotated: list[str] = []
    for index, line in enumerate(lines):
        annotated.extend(inserts.get(index, []))
        annotated.append(line)
    return "\n".join(annotated) + ("\n" if source.endswith("\n") else "")


def generate_block_comments(blocks: list[BlockInfo], model_choice: ModelChoice = "base") -> tuple[dict[int, str], list[str]]:
    comments: dict[int, str] = {}
    notes: list[str] = []

    for block in blocks:
        try:
            comments[block.lineno] = generate_comment(block.kind, block.name, block.source, variant=model_choice)
        except ModelUnavailableError:
            raise
        except Exception as exc:
            comments[block.lineno] = f"# TODO(model): Could not explain {block.kind} `{block.name}`."
            notes.append(f"{block.name}: model generation failed ({type(exc).__name__}).")
    return comments, notes


def annotate_python(source: str, model_choice: ModelChoice = "base") -> AnnotationResult:
    source = source.strip("\ufeff")
    if not source.strip():
        return AnnotationResult("", "", "Paste a Python file to annotate.", False, 0)

    try:
        original_tree = parse_python(source)
    except SyntaxError as exc:
        return AnnotationResult("", "", f"Syntax error on line {exc.lineno}: {exc.msg}", False, 0)

    blocks = collect_blocks(original_tree, source)
    summary, summary_notes = generate_summary(source, blocks, model_choice)
    if not blocks:
        status = "Parsed successfully. No classes or functions to annotate."
        if summary_notes:
            status += "\n" + "\n".join(summary_notes)
        return AnnotationResult(summary, source, status, True, 0, tuple(summary_notes))

    comments, generation_notes = generate_block_comments(blocks, model_choice)
    generation_notes = summary_notes + generation_notes
    annotated = insert_comments(source, comments)

    try:
        same_ast = semantic_ast_dump(source) == semantic_ast_dump(annotated)
    except SyntaxError as exc:
        return AnnotationResult(
            "",
            annotated,
            f"Generated annotation failed to parse on line {exc.lineno}: {exc.msg}",
            False,
            len(blocks),
            tuple(generation_notes),
        )

    status = (
        f"Validated: parsed {len(blocks)} block(s), inserted {model_choice} model comments, semantic AST unchanged."
        if same_ast
        else "Rejected: annotation changed the semantic AST."
    )
    if generation_notes:
        status += "\n" + "\n".join(generation_notes)

    return AnnotationResult(
        summary,
        annotated if same_ast else source,
        status,
        same_ast,
        len(blocks),
        tuple(generation_notes),
    )