ingest/parse_ast.py

"""
ast_parser.py
-------------
Parses Python source files using the built-in `ast` module to extract
structured metadata: classes, functions, parameters, return types,
docstrings, and call relationships (for cross-module analysis).

Uses `rich` for all terminal output.
Depends on file_walker.walk_codebase() for file discovery.
"""

import ast
import sys
from pathlib import Path
from dataclasses import dataclass, field

from rich.console import Console
from rich.tree import Tree
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich import box

# Import from file_walker (must be in same directory or on PYTHONPATH)
from ingest.walk_files import walk_codebase, group_by_module

console = Console()


# ── Data Models ───────────────────────────────────────────────────────────────

@dataclass
class ParameterInfo:
    name: str
    annotation: str | None = None
    default: str | None = None


@dataclass
class FunctionInfo:
    name: str
    parameters: list[ParameterInfo] = field(default_factory=list)
    return_type: str | None = None
    docstring: str | None = None
    calls: list[str] = field(default_factory=list)        # functions this calls
    lineno: int = 0
    is_method: bool = False
    source: str | None = None


@dataclass
class ClassInfo:
    name: str
    bases: list[str] = field(default_factory=list)
    docstring: str | None = None
    methods: list[FunctionInfo] = field(default_factory=list)
    lineno: int = 0
    source: str | None = None


@dataclass
class FileInfo:
    path: str
    relative: str
    module: str
    classes: list[ClassInfo] = field(default_factory=list)
    functions: list[FunctionInfo] = field(default_factory=list)   # top-level
    imports: list[str] = field(default_factory=list)
    docstring: str | None = None


# ── AST Helpers ───────────────────────────────────────────────────────────────

def _annotation_to_str(node) -> str | None:
    """Convert an AST annotation node to a readable string."""
    if node is None:
        return None
    try:
        return ast.unparse(node)
    except Exception:
        return None


def _extract_calls(func_node: ast.FunctionDef | ast.AsyncFunctionDef) -> list[str]:
    """Extract all function/method call names from a function body."""
    calls = []
    for node in ast.walk(func_node):
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name):
                calls.append(node.func.id)
            elif isinstance(node.func, ast.Attribute):
                calls.append(f"{ast.unparse(node.func.value)}.{node.func.attr}")
    return list(dict.fromkeys(calls))  # deduplicate, preserve order

def _parse_function(node: ast.FunctionDef | ast.AsyncFunctionDef,
                        is_method: bool = False,
                        raw_source: str | None = None) -> FunctionInfo:
    """Parse a function or method AST node into a FunctionInfo."""
    params = []

    args = node.args
    all_args = args.posonlyargs + args.args + args.kwonlyargs
    defaults_offset = len(all_args) - len(args.defaults)

    for i, arg in enumerate(all_args):
        if is_method and arg.arg == "self":
            continue
        default_val = None
        default_idx = i - defaults_offset
        if default_idx >= 0 and default_idx < len(args.defaults):
            try:
                default_val = ast.unparse(args.defaults[default_idx])
            except Exception:
                default_val = "..."

        params.append(ParameterInfo(
            name=arg.arg,
            annotation=_annotation_to_str(arg.annotation),
            default=default_val,
        ))

    # *args
    if args.vararg:
        params.append(ParameterInfo(
            name=f"*{args.vararg.arg}",
            annotation=_annotation_to_str(args.vararg.annotation),
        ))

    # **kwargs
    if args.kwarg:
        params.append(ParameterInfo(
            name=f"**{args.kwarg.arg}",
            annotation=_annotation_to_str(args.kwarg.annotation),
        ))

    source_segment = None
    if raw_source is not None:
        source_segment = ast.get_source_segment(raw_source, node)

    return FunctionInfo(
        name=node.name,
        parameters=params,
        return_type=_annotation_to_str(node.returns),
        docstring=ast.get_docstring(node),
        calls=_extract_calls(node),
        lineno=node.lineno,
        is_method=is_method,
        source=source_segment,
    )

def _parse_class(node: ast.ClassDef, raw_source: str | None = None) -> ClassInfo:
    """Parse a class AST node into a ClassInfo."""
    bases = []
    for base in node.bases:
        try:
            bases.append(ast.unparse(base))
        except Exception:
            pass

    methods = []
    for item in node.body:
        if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
            methods.append(_parse_function(item, is_method=True, raw_source=raw_source))

    class_source = None
    if raw_source is not None:
        class_source = ast.get_source_segment(raw_source, node)

    return ClassInfo(
        name=node.name,
        bases=bases,
        docstring=ast.get_docstring(node),
        methods=methods,
        lineno=node.lineno,
        source=class_source,
    )


def parse_file(file_meta: dict) -> FileInfo | None:
    """
    Parse a single Python source file and return a FileInfo.
    Returns None if the file cannot be parsed (syntax errors, encoding issues).

    Args:
        file_meta: A dict from walk_codebase() with keys: path, relative, module.
    """
    if file_meta["extension"] != ".py":
        return None  # AST parser only supports Python for now

    try:
        source = Path(file_meta["path"]).read_text(encoding="utf-8", errors="ignore")
        tree = ast.parse(source)
    except SyntaxError as e:
        console.print(f"[yellow]⚠ Syntax error in {file_meta['relative']}: {e}[/yellow]")
        return None
    except Exception as e:
        console.print(f"[yellow]⚠ Could not parse {file_meta['relative']}: {e}[/yellow]")
        return None

    info = FileInfo(
        path=file_meta["path"],
        relative=file_meta["relative"],
        module=file_meta["module"],
        docstring=ast.get_docstring(tree),
    )

    # Extract imports
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                info.imports.append(alias.name)
        elif isinstance(node, ast.ImportFrom):
            mod = node.module or ""
            for alias in node.names:
                info.imports.append(f"{mod}.{alias.name}")

    # Extract top-level classes and functions
    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            info.classes.append(_parse_class(node, raw_source=source))
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            info.functions.append(_parse_function(node, is_method=False, raw_source=source))

    return info


def parse_codebase(root_path: str) -> list[FileInfo]:
    """
    Walk and parse an entire codebase, returning FileInfo for every
    parseable Python file.

    Args:
        root_path: Absolute path to the codebase root.

    Returns:
        List of FileInfo objects (one per parsed file).
    """
    files = walk_codebase(root_path)
    py_files = [f for f in files if f["extension"] == ".py"]

    results = []
    with console.status(f"[bold cyan]Parsing {len(py_files)} Python files...[/bold cyan]"):
        for f in py_files:
            info = parse_file(f)
            if info:
                results.append(info)

    return results


# ── Rich Renderers ────────────────────────────────────────────────────────────

def render_file_tree(parsed_files: list[FileInfo]) -> None:
    """Render a rich tree of the parsed codebase structure."""
    root_tree = Tree(
        "[bold bright_cyan]📦 Codebase Structure[/bold bright_cyan]",
        guide_style="dim cyan",
    )

    # Group by module
    modules: dict[str, list[FileInfo]] = {}
    for f in parsed_files:
        modules.setdefault(f.module, []).append(f)

    for module_name, module_files in sorted(modules.items()):
        module_branch = root_tree.add(
            f"[bold yellow]📁 {module_name}[/bold yellow]"
        )
        for file_info in module_files:
            file_label = f"[green]{Path(file_info.relative).name}[/green]"
            file_branch = module_branch.add(file_label)

            for cls in file_info.classes:
                cls_branch = file_branch.add(
                    f"[bold magenta]🔷 {cls.name}[/bold magenta]"
                )
                for method in cls.methods:
                    cls_branch.add(f"[cyan]  ƒ {method.name}()[/cyan]")

            for func in file_info.functions:
                file_branch.add(f"[blue]  ƒ {func.name}()[/blue]")

    console.print(root_tree)


def render_function(func: FunctionInfo, class_name: str | None = None) -> None:
    """Render a detailed view of a single function/method."""
    title = f"{'Method' if func.is_method else 'Function'}: "
    title += f"{class_name}.{func.name}" if class_name else func.name

    # Signature table
    param_table = Table(
        "Parameter", "Type", "Default",
        box=box.SIMPLE_HEAD,
        style="dim",
        header_style="bold cyan",
        show_edge=False,
    )
    for p in func.parameters:
        param_table.add_row(
            p.name,
            p.annotation or "[dim]any[/dim]",
            p.default or "[dim]—[/dim]",
        )

    return_str = func.return_type or "[dim]None / untyped[/dim]"
    doc_str = func.docstring or "[dim]No docstring.[/dim]"
    calls_str = ", ".join(func.calls[:10]) if func.calls else "[dim]none[/dim]"

    content = Text()
    content.append("📄 Docstring\n", style="bold")
    content.append(f"{doc_str}\n\n")
    content.append("↩  Return type\n", style="bold")
    content.append(f"{return_str}\n\n")
    content.append("📞 Calls\n", style="bold")
    content.append(calls_str)

    console.print(Panel(content, title=f"[bold white]{title}[/bold white]",
                        border_style="cyan"))
    if func.parameters:
        console.print(param_table)
    else:
        console.print("  [dim]No parameters.[/dim]\n")


def render_class(cls: ClassInfo) -> None:
    """Render a detailed view of a class and all its methods."""
    bases_str = ", ".join(cls.bases) if cls.bases else "object"
    doc_str = cls.docstring or "[dim]No docstring.[/dim]"

    header = Text()
    header.append(f"class {cls.name}", style="bold magenta")
    header.append(f"({bases_str})", style="dim")

    console.print(Panel(
        f"[bold]Docstring:[/bold]\n{doc_str}",
        title=str(header),
        border_style="magenta",
    ))

    for method in cls.methods:
        render_function(method, class_name=cls.name)


def render_summary(parsed_files: list[FileInfo]) -> None:
    """Render a high-level summary table of the parsed codebase."""
    table = Table(
        "Module", "File", "Classes", "Functions", "Imports",
        box=box.ROUNDED,
        header_style="bold bright_cyan",
        border_style="cyan",
        show_lines=True,
    )

    for f in parsed_files:
        table.add_row(
            f.module,
            Path(f.relative).name,
            str(len(f.classes)),
            str(len(f.functions)),
            str(len(f.imports)),
        )

    total_classes   = sum(len(f.classes)   for f in parsed_files)
    total_functions = sum(len(f.functions) for f in parsed_files)
    total_methods   = sum(len(c.methods)   for f in parsed_files for c in f.classes)

    console.print(table)
    console.print(
        f"\n[bold]Totals:[/bold] "
        f"[cyan]{len(parsed_files)}[/cyan] files · "
        f"[magenta]{total_classes}[/magenta] classes · "
        f"[blue]{total_functions}[/blue] top-level functions · "
        f"[green]{total_methods}[/green] methods\n"
    )


# ── Entry Point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    path = sys.argv[1] if len(sys.argv) > 1 else "."

    console.print(f"\n[bold cyan]🔍 Parsing codebase:[/bold cyan] {path}\n")

    try:
        parsed = parse_codebase(path)
    except (FileNotFoundError, NotADirectoryError) as e:
        console.print(f"[red]❌ Error: {e}[/red]")
        sys.exit(1)

    if not parsed:
        console.print("[yellow]No Python files found or parsed.[/yellow]")
        sys.exit(0)

    console.rule("[bold cyan]Summary[/bold cyan]")
    render_summary(parsed)

    console.rule("[bold cyan]Codebase Tree[/bold cyan]")
    render_file_tree(parsed)

    # Demo: render first class found
    for f in parsed:
        if f.classes:
            console.rule(f"[bold cyan]Sample Class — {f.classes[0].name}[/bold cyan]")
            render_class(f.classes[0])
            break

    # Demo: render first top-level function found
    for f in parsed:
        if f.functions:
            console.rule(f"[bold cyan]Sample Function — {f.functions[0].name}[/bold cyan]")
            render_function(f.functions[0])
            break