parser.py

"""Parse Python source files into a structured representation.

For each file we extract: imports, top-level functions, classes (with methods),
module-level constants, and per-symbol structural facts — signatures,
decorators, base classes, instance attributes — that are derived directly from
the AST. These are the ground truth; docstrings are captured separately as
"documented intent" and should not be conflated with what the code actually
does.

Call resolution lives in graph.py; this module records calls as the dotted
name they were written with (`foo`, `os.path.join`, `self.bar`).
"""

import ast
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional


@dataclass
class Import:
    module: str                  # module being imported, e.g. "os.path" or ".sibling"
    name: Optional[str]          # symbol name when "from x import y", else None
    alias: str                   # local name (the asname if present, else module/name)
    lineno: int = 0              # line of the import statement


@dataclass
class Param:
    name: str
    annotation: Optional[str] = None
    default: Optional[str] = None
    kind: str = "pos"            # "pos" | "vararg" | "kwarg" | "kwonly"


@dataclass
class CallSite:
    name: str        # dotted name as written, e.g. "self.do_thing" or "os.path.join"
    lineno: int      # line in the source file where the call occurs


@dataclass
class FunctionInfo:
    name: str
    qualname: str
    lineno: int
    end_lineno: int = 0
    params: list[Param] = field(default_factory=list)
    returns: Optional[str] = None
    decorators: list[str] = field(default_factory=list)
    is_async: bool = False
    calls: list[CallSite] = field(default_factory=list)
    local_attrs_written: list[str] = field(default_factory=list)
    doc: Optional[str] = None
    source: str = ""                                              # def block sliced from the file


@dataclass
class ClassInfo:
    name: str
    lineno: int
    end_lineno: int = 0
    bases: list[str] = field(default_factory=list)
    decorators: list[str] = field(default_factory=list)
    class_attrs: list[str] = field(default_factory=list)
    instance_attrs: list[str] = field(default_factory=list)
    methods: list[FunctionInfo] = field(default_factory=list)
    doc: Optional[str] = None
    source: str = ""


@dataclass
class FileInfo:
    path: Path
    relpath: str
    module: str
    imports: list[Import] = field(default_factory=list)
    functions: list[FunctionInfo] = field(default_factory=list)
    classes: list[ClassInfo] = field(default_factory=list)
    constants: list[str] = field(default_factory=list)
    doc: Optional[str] = None
    parse_error: Optional[str] = None
    header_source: str = ""                                       # first ~20 lines, the file's elevator pitch
    total_lines: int = 0


# ---------- low-level helpers ----------

def _dotted_name(node: ast.AST) -> Optional[str]:
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        base = _dotted_name(node.value)
        if base is None:
            return None
        return f"{base}.{node.attr}"
    return None


def _unparse(node: Optional[ast.AST]) -> Optional[str]:
    if node is None:
        return None
    try:
        return ast.unparse(node)
    except Exception:
        return None


def _extract_calls(body: list[ast.stmt]) -> list[CallSite]:
    calls: list[CallSite] = []
    for node in body:
        for child in ast.walk(node):
            if isinstance(child, ast.Call):
                name = _dotted_name(child.func)
                if name:
                    calls.append(CallSite(name=name, lineno=child.lineno))
    return calls


def _slice_source(source_lines: list[str], start: int, end: int,
                  max_lines: int = 80) -> str:
    """Slice a 1-based inclusive line range and prefix each line with its
    line number so the output is reader-friendly. Truncates the middle when
    the span exceeds `max_lines`."""
    if not source_lines or start < 1 or end < start:
        return ""
    end = min(end, len(source_lines))
    numbered = list(enumerate(source_lines[start - 1:end], start=start))  # [(N, line)]
    width = len(str(end))
    fmt = lambda n, ln: f"{str(n).rjust(width)}: {ln}"

    if len(numbered) <= max_lines:
        return "\n".join(fmt(n, ln) for n, ln in numbered)

    head = numbered[: max_lines - 6]
    tail = numbered[-5:]
    elided = len(numbered) - len(head) - len(tail)
    head_out = [fmt(n, ln) for n, ln in head]
    tail_out = [fmt(n, ln) for n, ln in tail]
    gap = " " * width + f"    # … ({elided} more lines elided) …"
    return "\n".join(head_out + [gap] + tail_out)


def _header_source(source_lines: list[str], max_lines: int = 20) -> str:
    """Capture the file's elevator pitch — first ~20 lines, numbered."""
    if not source_lines:
        return ""
    end = min(max_lines, len(source_lines))
    return _slice_source(source_lines, 1, end, max_lines=max_lines + 1)


def _extract_imports(tree: ast.Module) -> list[Import]:
    imports: list[Import] = []
    for node in ast.iter_child_nodes(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                local = alias.asname or alias.name.split(".")[0]
                imports.append(Import(module=alias.name, name=None, alias=local,
                                      lineno=node.lineno))
        elif isinstance(node, ast.ImportFrom):
            module = node.module or ""
            prefix = "." * (node.level or 0)
            full = f"{prefix}{module}" if prefix else module
            for alias in node.names:
                local = alias.asname or alias.name
                imports.append(Import(module=full, name=alias.name, alias=local,
                                      lineno=node.lineno))
    return imports


def _module_name_from_relpath(relpath: str) -> str:
    parts = list(Path(relpath).with_suffix("").parts)
    if parts and parts[-1] == "__init__":
        parts.pop()
    return ".".join(parts)


def _extract_params(args: ast.arguments) -> list[Param]:
    """Flatten ast.arguments into a list of Param facts."""
    out: list[Param] = []
    defaults = list(args.defaults)
    posonly = list(getattr(args, "posonlyargs", []))
    pos = list(args.args)
    # Defaults align to the tail of (posonly + pos) combined
    all_pos = posonly + pos
    n_with_default = len(defaults)
    default_offset = len(all_pos) - n_with_default
    for i, a in enumerate(all_pos):
        d = defaults[i - default_offset] if i >= default_offset else None
        out.append(Param(name=a.arg, annotation=_unparse(a.annotation),
                         default=_unparse(d), kind="pos"))
    if args.vararg:
        out.append(Param(name=args.vararg.arg, annotation=_unparse(args.vararg.annotation),
                         kind="vararg"))
    kw_defaults = args.kw_defaults
    for i, a in enumerate(args.kwonlyargs):
        d = kw_defaults[i] if i < len(kw_defaults) else None
        out.append(Param(name=a.arg, annotation=_unparse(a.annotation),
                         default=_unparse(d), kind="kwonly"))
    if args.kwarg:
        out.append(Param(name=args.kwarg.arg, annotation=_unparse(args.kwarg.annotation),
                         kind="kwarg"))
    return out


def _self_attrs_written(body: list[ast.stmt]) -> list[str]:
    """Find all `self.X = ...` and `self.X: T = ...` writes within a method body."""
    found: list[str] = []
    seen: set[str] = set()
    for node in body:
        for child in ast.walk(node):
            targets = []
            if isinstance(child, ast.Assign):
                targets = child.targets
            elif isinstance(child, ast.AnnAssign):
                targets = [child.target]
            elif isinstance(child, ast.AugAssign):
                targets = [child.target]
            for t in targets:
                if isinstance(t, ast.Attribute) and isinstance(t.value, ast.Name) \
                        and t.value.id == "self":
                    if t.attr not in seen:
                        seen.add(t.attr)
                        found.append(t.attr)
                # tuple unpacking: (self.a, self.b) = ...
                if isinstance(t, ast.Tuple):
                    for el in t.elts:
                        if isinstance(el, ast.Attribute) and isinstance(el.value, ast.Name) \
                                and el.value.id == "self":
                            if el.attr not in seen:
                                seen.add(el.attr)
                                found.append(el.attr)
    return found


def _decorator_strings(decos: list[ast.expr]) -> list[str]:
    out: list[str] = []
    for d in decos:
        s = _unparse(d)
        if s:
            out.append(s)
    return out


def _function_info(node: ast.FunctionDef | ast.AsyncFunctionDef, qualname: str,
                   source_lines: list[str]) -> FunctionInfo:
    end = node.end_lineno or node.lineno
    return FunctionInfo(
        name=node.name,
        qualname=qualname,
        lineno=node.lineno,
        end_lineno=end,
        params=_extract_params(node.args),
        returns=_unparse(node.returns),
        decorators=_decorator_strings(node.decorator_list),
        is_async=isinstance(node, ast.AsyncFunctionDef),
        calls=_extract_calls(node.body),
        local_attrs_written=_self_attrs_written(node.body),
        doc=ast.get_docstring(node, clean=True),
        source=_slice_source(source_lines, node.lineno, end),
    )


def _class_attrs(body: list[ast.stmt]) -> list[str]:
    """Names assigned at class scope (excluding methods/nested classes)."""
    out: list[str] = []
    seen: set[str] = set()
    for node in body:
        if isinstance(node, ast.Assign):
            for t in node.targets:
                if isinstance(t, ast.Name) and t.id not in seen:
                    seen.add(t.id)
                    out.append(t.id)
        elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
            if node.target.id not in seen:
                seen.add(node.target.id)
                out.append(node.target.id)
    return out


def _module_constants(tree: ast.Module) -> list[str]:
    """Top-level UPPERCASE / SCREAMING_SNAKE constants."""
    out: list[str] = []
    seen: set[str] = set()
    for node in ast.iter_child_nodes(tree):
        names: list[str] = []
        if isinstance(node, ast.Assign):
            for t in node.targets:
                if isinstance(t, ast.Name):
                    names.append(t.id)
        elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
            names.append(node.target.id)
        for n in names:
            if n.isupper() and not n.startswith("_") and n not in seen:
                seen.add(n)
                out.append(n)
    return out


# ---------- main entry ----------

def parse_file(path: Path, root: Path) -> FileInfo:
    relpath = str(path.resolve().relative_to(root.resolve())).replace("\\", "/")
    info = FileInfo(path=path, relpath=relpath, module=_module_name_from_relpath(relpath))

    try:
        source = path.read_text(encoding="utf-8")
    except (UnicodeDecodeError, OSError) as e:
        info.parse_error = f"read failed: {e}"
        return info

    try:
        tree = ast.parse(source, filename=str(path))
    except SyntaxError as e:
        info.parse_error = f"syntax error at line {e.lineno}: {e.msg}"
        return info

    info.imports = _extract_imports(tree)
    info.doc = ast.get_docstring(tree, clean=True)
    info.constants = _module_constants(tree)

    source_lines = source.splitlines()
    info.total_lines = len(source_lines)
    info.header_source = _header_source(source_lines)

    for node in ast.iter_child_nodes(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            info.functions.append(_function_info(node, node.name, source_lines))
        elif isinstance(node, ast.ClassDef):
            cls_end = node.end_lineno or node.lineno
            cls = ClassInfo(
                name=node.name,
                lineno=node.lineno,
                end_lineno=cls_end,
                bases=[s for s in (_unparse(b) for b in node.bases) if s],
                decorators=_decorator_strings(node.decorator_list),
                class_attrs=_class_attrs(node.body),
                doc=ast.get_docstring(node, clean=True),
                source=_slice_source(source_lines, node.lineno, cls_end),
            )
            instance_seen: set[str] = set()
            for sub in node.body:
                if isinstance(sub, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    meth = _function_info(sub, f"{node.name}.{sub.name}", source_lines)
                    cls.methods.append(meth)
                    for a in meth.local_attrs_written:
                        if a not in instance_seen:
                            instance_seen.add(a)
                            cls.instance_attrs.append(a)
            info.classes.append(cls)
    return info


def parse_codebase(root: str | Path, files) -> list[FileInfo]:
    root_path = Path(root).resolve()
    return [parse_file(Path(f), root_path) for f in files]


# ---------- fact-based summarization ----------

def format_signature(fn: FunctionInfo) -> str:
    """Render a Python-style signature: `name(a: int, b='x') -> T`."""
    parts: list[str] = []
    saw_kwonly_marker = False
    for p in fn.params:
        if p.kind == "vararg":
            parts.append("*" + p.name + (f": {p.annotation}" if p.annotation else ""))
        elif p.kind == "kwarg":
            parts.append("**" + p.name + (f": {p.annotation}" if p.annotation else ""))
        elif p.kind == "kwonly":
            if not saw_kwonly_marker and not any(pp.kind == "vararg" for pp in fn.params):
                parts.append("*")
                saw_kwonly_marker = True
            s = p.name
            if p.annotation: s += f": {p.annotation}"
            if p.default is not None: s += f"={p.default}"
            parts.append(s)
        else:
            s = p.name
            if p.annotation: s += f": {p.annotation}"
            if p.default is not None: s += f"={p.default}"
            parts.append(s)
    sig = f"{fn.name}({', '.join(parts)})"
    if fn.returns:
        sig += f" -> {fn.returns}"
    if fn.is_async:
        sig = "async " + sig
    return sig


def summarize_function(fn: FunctionInfo) -> str:
    """A short, factual sentence about what `fn` is, derived purely from AST facts."""
    sig = format_signature(fn)
    bits: list[str] = []
    if fn.decorators:
        # Highlight a few semantic decorators
        decos = ", ".join("@" + d for d in fn.decorators)
        bits.append(decos)
    n_calls = len({c.name for c in fn.calls})
    if n_calls:
        bits.append(f"calls {n_calls} symbol{'s' if n_calls != 1 else ''}")
    if fn.local_attrs_written:
        bits.append(f"writes self.{{{', '.join(fn.local_attrs_written[:4])}}}"
                    + ("…" if len(fn.local_attrs_written) > 4 else ""))
    tail = " · ".join(bits)
    return f"{sig}" + (f"  ·  {tail}" if tail else "")


def summarize_class(cls: ClassInfo) -> str:
    parts = [f"class {cls.name}"]
    if cls.bases:
        parts.append("(" + ", ".join(cls.bases) + ")")
    head = "".join(parts)
    bits: list[str] = []
    if cls.decorators:
        bits.append(", ".join("@" + d for d in cls.decorators))
    bits.append(f"{len(cls.methods)} method{'s' if len(cls.methods) != 1 else ''}")
    if cls.instance_attrs:
        n = len(cls.instance_attrs)
        bits.append(f"{n} instance attr{'s' if n != 1 else ''}")
    return head + "  ·  " + " · ".join(bits)


def summarize_file(f: FileInfo) -> str:
    bits: list[str] = []
    bits.append(f"{len(f.classes)} class{'es' if len(f.classes) != 1 else ''}")
    bits.append(f"{len(f.functions)} function{'s' if len(f.functions) != 1 else ''}")
    bits.append(f"{len(f.imports)} import{'s' if len(f.imports) != 1 else ''}")
    if f.constants:
        bits.append(f"{len(f.constants)} constant{'s' if len(f.constants) != 1 else ''}")
    return " · ".join(bits)


# ---------- "actual intent" — fact-derived narration ----------
#
# These compose a short, factual sentence about what the code DOES, using only
# the structural facts already extracted. Each clause is dropped when its data
# is missing, so the output stays honest (no fabricated phrases).


def _param_for_intent(p: Param) -> str:
    s = p.name
    if p.annotation:
        s += f": {p.annotation}"
    return s


def derive_function_intent(fn: FunctionInfo) -> str:
    parts: list[str] = []

    visible = [p for p in fn.params if p.name not in ("self", "cls")]
    if visible:
        items = [_param_for_intent(p) for p in visible[:4]]
        more = f" (+{len(visible) - 4} more)" if len(visible) > 4 else ""
        parts.append(f"Takes {', '.join(items)}{more}.")

    if fn.returns:
        parts.append(f"Returns {fn.returns}.")

    unique_calls: list[str] = []
    seen: set[str] = set()
    for c in fn.calls:
        if c.name not in seen:
            seen.add(c.name)
            unique_calls.append(c.name)
    if unique_calls:
        items = unique_calls[:5]
        more = f" (+{len(unique_calls) - 5} more)" if len(unique_calls) > 5 else ""
        parts.append(f"Composes: {', '.join(items)}{more}.")

    if fn.local_attrs_written:
        attrs = fn.local_attrs_written[:4]
        more = f" (+{len(fn.local_attrs_written) - 4})" if len(fn.local_attrs_written) > 4 else ""
        parts.append(f"Mutates: self.{{{', '.join(attrs)}{more}}}.")

    if fn.is_async:
        parts.insert(0, "Async coroutine.")

    return " ".join(parts) if parts else "(no observable behavior derived from AST)"


def derive_class_intent(cls: ClassInfo) -> str:
    parts: list[str] = []
    if cls.bases:
        parts.append(f"Subclasses {', '.join(cls.bases)}.")
    n = len(cls.methods)
    if n:
        parts.append(f"Defines {n} method{'s' if n != 1 else ''}.")
    if cls.instance_attrs:
        attrs = cls.instance_attrs[:5]
        more = f" (+{len(cls.instance_attrs) - 5} more)" if len(cls.instance_attrs) > 5 else ""
        parts.append(f"Holds instance state: {', '.join(attrs)}{more}.")
    if cls.class_attrs:
        ca = cls.class_attrs[:4]
        parts.append(f"Class-level attrs: {', '.join(ca)}.")
    return " ".join(parts) if parts else "(empty class body)"


def derive_file_intent(f: FileInfo) -> str:
    parts: list[str] = []

    if f.classes:
        names = [c.name for c in f.classes[:4]]
        more = f" (+{len(f.classes) - 4} more)" if len(f.classes) > 4 else ""
        parts.append(f"Defines {', '.join(names)}{more}.")

    if f.functions:
        names = [fn.name for fn in f.functions[:4]]
        more = f" (+{len(f.functions) - 4} more)" if len(f.functions) > 4 else ""
        parts.append(f"Exports {', '.join(names)}{more}.")

    # Top external libs by occurrence in this file's imports
    counts: dict[str, int] = {}
    for imp in f.imports:
        top = imp.module.split(".")[0].lstrip(".")
        if top:
            counts[top] = counts.get(top, 0) + 1
    ordered = sorted(counts.items(), key=lambda kv: -kv[1])
    if ordered:
        libs = [name for name, _ in ordered[:4]]
        parts.append(f"Pulls in {', '.join(libs)}.")

    if f.constants:
        parts.append(f"Module constants: {', '.join(f.constants[:4])}.")

    return " ".join(parts) if parts else "(empty module)"


if __name__ == "__main__":
    import sys
    from walker import walk_python_files

    target = sys.argv[1] if len(sys.argv) > 1 else "."
    root = Path(target).resolve()
    files = list(walk_python_files(root))
    parsed = parse_codebase(root, files)

    failed = [f for f in parsed if f.parse_error]
    print(f"parsed {len(parsed)} files ({len(failed)} failed)")

    if parsed:
        sample = next((f for f in parsed if f.functions or f.classes), parsed[0])
        print(f"\nsample: {sample.relpath}  ·  {summarize_file(sample)}")
        for fn in sample.functions[:3]:
            print("  fn", summarize_function(fn))
        for cls in sample.classes[:2]:
            print("  ", summarize_class(cls))
            for m in cls.methods[:3]:
                print("    ", summarize_function(m))