DocForge2/src/docforge/python_logic.py

"""Deterministic, function-scoped Python control-flow extraction.

The analyzer parses source as data. It never imports or executes project code.
Its projections intentionally remain separate from DocForge's primary graph.
"""

from __future__ import annotations

import ast
import hashlib
from collections.abc import Iterable
from dataclasses import dataclass

from .errors import DocForgeError
from .models import LogicEdge, LogicNode, LogicProjection

FunctionNode = ast.FunctionDef | ast.AsyncFunctionDef


@dataclass(frozen=True)
class PythonLogicOwner:
    """One primary graph function or method that should receive a logic projection."""

    owner_node_id: str
    qualified_name: str
    line: int


@dataclass(frozen=True)
class _Tail:
    source_id: str
    relation: str = "next"
    label: str | None = None


@dataclass(frozen=True)
class _Condition:
    entry_id: str
    when_true: tuple[_Tail, ...]
    when_false: tuple[_Tail, ...]


@dataclass(frozen=True)
class _Loop:
    continue_id: str
    break_id: str


def analyze_python_source(
    source: str,
    *,
    source_id: str,
    owners: Iterable[PythonLogicOwner],
    filename: str = "<python-source>",
    max_nodes_per_function: int = 2_000,
) -> tuple[LogicProjection, ...]:
    """Build ordered control-flow projections for explicitly owned Python functions."""

    if max_nodes_per_function < 2:
        raise ValueError("max_nodes_per_function must allow entry and exit nodes")
    try:
        tree = ast.parse(source, filename=filename)
    except SyntaxError as error:
        raise DocForgeError(
            "invalid_logic_source",
            "Python source cannot be parsed for logic analysis",
            source=filename,
            line=error.lineno,
        ) from error
    definitions = _function_definitions(tree)
    requested = tuple(sorted(owners, key=lambda item: item.owner_node_id))
    if len({owner.owner_node_id for owner in requested}) != len(requested):
        raise DocForgeError("invalid_logic_owner", "Logic owner IDs must be unique")
    projections: list[LogicProjection] = []
    for owner in requested:
        function = definitions.get((owner.qualified_name, owner.line))
        if function is None:
            raise DocForgeError(
                "missing_logic_owner",
                "A requested Python logic owner was not found in its source",
                owner_node_id=owner.owner_node_id,
                qualified_name=owner.qualified_name,
                line=owner.line,
            )
        projections.append(
            _FunctionLogicBuilder(
                source_id=source_id,
                owner_node_id=owner.owner_node_id,
                function=function,
                max_nodes=max_nodes_per_function,
            ).build()
        )
    return tuple(projections)


def _function_definitions(tree: ast.Module) -> dict[tuple[str, int], FunctionNode]:
    result: dict[tuple[str, int], FunctionNode] = {}

    class DefinitionVisitor(ast.NodeVisitor):
        def __init__(self) -> None:
            self.parents: tuple[str, ...] = ()

        def _visit_scope(self, name: str, body: list[ast.stmt]) -> None:
            previous = self.parents
            self.parents = (*previous, name)
            for statement in body:
                self.visit(statement)
            self.parents = previous

        def visit_ClassDef(self, node: ast.ClassDef) -> None:
            self._visit_scope(node.name, node.body)

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            qualified_name = ".".join((*self.parents, node.name))
            result[(qualified_name, node.lineno)] = node
            self._visit_scope(node.name, node.body)

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            qualified_name = ".".join((*self.parents, node.name))
            result[(qualified_name, node.lineno)] = node
            self._visit_scope(node.name, node.body)

    DefinitionVisitor().visit(tree)
    return result


class _FunctionLogicBuilder:
    def __init__(
        self,
        *,
        source_id: str,
        owner_node_id: str,
        function: FunctionNode,
        max_nodes: int,
    ) -> None:
        self.source_id = source_id
        self.owner_node_id = owner_node_id
        self.function = function
        self.max_nodes = max_nodes
        self.nodes: list[LogicNode] = []
        self.edges: list[LogicEdge] = []
        self._edge_ordinals: dict[str, int] = {}
        self._sequence = 0
        self._owner_digest = hashlib.sha256(owner_node_id.encode()).hexdigest()[:12]
        self.entry_id = self._node("entry", f"Enter {function.name}", function)
        self.exit_id = self._node("exit", f"Exit {function.name}", function)

    def build(self) -> LogicProjection:
        incoming = (_Tail(self.entry_id),)
        body = list(self.function.body)
        if body and _is_docstring(body[0]):
            body = body[1:]
        tails = self._statements(body, incoming, loop=None)
        self._connect(tails, self.exit_id)
        if not self._has_incoming(self.exit_id):
            self._edge(self.entry_id, "next", self.exit_id, "END")
        return LogicProjection(
            owner_node_id=self.owner_node_id,
            source_id=self.source_id,
            nodes=tuple(sorted(self.nodes, key=lambda node: node.logic_id)),
            edges=tuple(
                sorted(
                    self.edges,
                    key=lambda edge: (
                        edge.source_id,
                        edge.ordinal,
                        edge.relation,
                        edge.target_id,
                        edge.label or "",
                    ),
                )
            ),
        )

    def _statements(
        self,
        statements: list[ast.stmt],
        incoming: tuple[_Tail, ...],
        *,
        loop: _Loop | None,
    ) -> tuple[_Tail, ...]:
        tails = incoming
        for statement in statements:
            if not tails:
                break
            tails = self._statement(statement, tails, loop=loop)
        return tails

    def _statement(
        self,
        statement: ast.stmt,
        incoming: tuple[_Tail, ...],
        *,
        loop: _Loop | None,
    ) -> tuple[_Tail, ...]:
        if isinstance(statement, ast.If):
            return self._if(statement, incoming, loop=loop)
        if isinstance(statement, (ast.While,)):
            return self._while(statement, incoming)
        if isinstance(statement, (ast.For, ast.AsyncFor)):
            return self._for(statement, incoming)
        if isinstance(statement, ast.Match):
            return self._match(statement, incoming, loop=loop)
        if isinstance(statement, (ast.Try, ast.TryStar)):
            return self._try(statement, incoming, loop=loop)
        if isinstance(statement, (ast.With, ast.AsyncWith)):
            label = f"{'async ' if isinstance(statement, ast.AsyncWith) else ''}with "
            label += ", ".join(_expression(item.context_expr) for item in statement.items)
            node_id = self._node("action", label, statement)
            self._connect(incoming, node_id)
            return self._statements(statement.body, (_Tail(node_id),), loop=loop)
        if isinstance(statement, ast.Return):
            label = (
                "return" if statement.value is None else f"return {_expression(statement.value)}"
            )
            node_id = self._node("return", label, statement)
            self._connect(incoming, node_id)
            self._edge(node_id, "return", self.exit_id, "RETURN")
            return ()
        if isinstance(statement, ast.Raise):
            label = "raise" if statement.exc is None else f"raise {_expression(statement.exc)}"
            node_id = self._node("raise", label, statement)
            self._connect(incoming, node_id)
            self._edge(node_id, "raise", self.exit_id, "RAISE")
            return ()
        if isinstance(statement, ast.Break):
            node_id = self._node("break", "break", statement)
            self._connect(incoming, node_id)
            if loop is not None:
                self._edge(node_id, "break", loop.break_id, "BREAK")
            else:
                self._edge(node_id, "next", self.exit_id, "INVALID BREAK")
            return ()
        if isinstance(statement, ast.Continue):
            node_id = self._node("continue", "continue", statement)
            self._connect(incoming, node_id)
            if loop is not None:
                self._edge(node_id, "continue", loop.continue_id, "CONTINUE")
            else:
                self._edge(node_id, "next", self.exit_id, "INVALID CONTINUE")
            return ()
        if isinstance(statement, ast.Assert):
            condition = self._condition(statement.test, incoming)
            failure = self._node(
                "raise",
                "AssertionError"
                if statement.msg is None
                else f"AssertionError: {_expression(statement.msg)}",
                statement,
            )
            self._connect(condition.when_false, failure)
            self._edge(failure, "raise", self.exit_id, "RAISE")
            return condition.when_true

        kind = "call" if _contains_runtime_call(statement) else "action"
        node_id = self._node(kind, _statement_label(statement), statement)
        self._connect(incoming, node_id)
        return (_Tail(node_id),)

    def _if(
        self,
        statement: ast.If,
        incoming: tuple[_Tail, ...],
        *,
        loop: _Loop | None,
    ) -> tuple[_Tail, ...]:
        condition = self._condition(statement.test, incoming)
        body_tails = self._statements(statement.body, condition.when_true, loop=loop)
        else_tails = (
            self._statements(statement.orelse, condition.when_false, loop=loop)
            if statement.orelse
            else condition.when_false
        )
        return self._converge(
            "Decision convergence",
            (*body_tails, *else_tails),
            statement,
        )

    def _while(self, statement: ast.While, incoming: tuple[_Tail, ...]) -> tuple[_Tail, ...]:
        condition = self._condition(statement.test, incoming)
        after_id = self._node("convergence", "Loop exit", statement)
        loop = _Loop(continue_id=condition.entry_id, break_id=after_id)
        body_tails = self._statements(statement.body, condition.when_true, loop=loop)
        for tail in body_tails:
            self._edge(tail.source_id, "loop", condition.entry_id, "LOOP")
        normal_tails = (
            self._statements(statement.orelse, condition.when_false, loop=None)
            if statement.orelse
            else condition.when_false
        )
        self._connect(normal_tails, after_id)
        return (_Tail(after_id),) if self._has_incoming(after_id) else ()

    def _for(
        self,
        statement: ast.For | ast.AsyncFor,
        incoming: tuple[_Tail, ...],
    ) -> tuple[_Tail, ...]:
        prefix = "async for" if isinstance(statement, ast.AsyncFor) else "for"
        loop_id = self._node(
            "loop",
            f"{prefix} {_expression(statement.target)} in {_expression(statement.iter)}",
            statement,
        )
        after_id = self._node("convergence", "Loop exit", statement)
        self._connect(incoming, loop_id)
        loop = _Loop(continue_id=loop_id, break_id=after_id)
        body_tails = self._statements(
            statement.body,
            (_Tail(loop_id, "when_true", "ITEM"),),
            loop=loop,
        )
        for tail in body_tails:
            self._edge(tail.source_id, "loop", loop_id, "NEXT ITEM")
        exhausted = (_Tail(loop_id, "when_false", "EXHAUSTED"),)
        normal_tails = (
            self._statements(statement.orelse, exhausted, loop=None)
            if statement.orelse
            else exhausted
        )
        self._connect(normal_tails, after_id)
        return (_Tail(after_id),) if self._has_incoming(after_id) else ()

    def _match(
        self,
        statement: ast.Match,
        incoming: tuple[_Tail, ...],
        *,
        loop: _Loop | None,
    ) -> tuple[_Tail, ...]:
        match_id = self._node("condition", f"match {_expression(statement.subject)}", statement)
        self._connect(incoming, match_id)
        pending: tuple[_Tail, ...] = (_Tail(match_id, "case", "CASE"),)
        completed: list[_Tail] = []
        for case in statement.cases:
            label = f"case {_expression(case.pattern)}"
            if case.guard is not None:
                label += f" if {_expression(case.guard)}"
            case_id = self._node("case", label, case.pattern)
            self._connect(pending, case_id)
            completed.extend(
                self._statements(
                    case.body,
                    (_Tail(case_id, "when_true", "MATCH"),),
                    loop=loop,
                )
            )
            pending = () if _is_catch_all(case) else (_Tail(case_id, "when_false", "NEXT CASE"),)
        return self._converge(
            "Case convergence",
            (*completed, *pending),
            statement,
        )

    def _try(
        self,
        statement: ast.Try | ast.TryStar,
        incoming: tuple[_Tail, ...],
        *,
        loop: _Loop | None,
    ) -> tuple[_Tail, ...]:
        try_id = self._node("try", "try", statement)
        self._connect(incoming, try_id)
        normal = self._statements(statement.body, (_Tail(try_id),), loop=loop)
        if statement.orelse:
            normal = self._statements(statement.orelse, normal, loop=loop)
        branches: list[_Tail] = list(normal)
        for handler in statement.handlers:
            exception = "Exception" if handler.type is None else _expression(handler.type)
            if handler.name:
                exception += f" as {handler.name}"
            handler_id = self._node("except", f"except {exception}", handler)
            self._edge(try_id, "exception", handler_id, f"EXCEPT {exception}")
            branches.extend(self._statements(handler.body, (_Tail(handler_id),), loop=loop))
        converged = self._converge(
            "Exception convergence",
            tuple(branches),
            statement,
        )
        if not statement.finalbody:
            return converged
        finally_id = self._node("finally", "finally", statement.finalbody[0])
        self._connect(converged, finally_id)
        return self._statements(statement.finalbody, (_Tail(finally_id),), loop=loop)

    def _condition(
        self,
        expression: ast.expr,
        incoming: tuple[_Tail, ...],
    ) -> _Condition:
        if isinstance(expression, ast.UnaryOp) and isinstance(expression.op, ast.Not):
            inner = self._condition(expression.operand, incoming)
            return _Condition(inner.entry_id, inner.when_false, inner.when_true)
        if isinstance(expression, ast.BoolOp) and expression.values:
            first = self._condition(expression.values[0], incoming)
            entry_id = first.entry_id
            if isinstance(expression.op, ast.And):
                when_true = first.when_true
                when_false = list(first.when_false)
                for value in expression.values[1:]:
                    next_condition = self._condition(value, when_true)
                    when_true = next_condition.when_true
                    when_false.extend(next_condition.when_false)
                return _Condition(entry_id, when_true, tuple(when_false))
            when_true = list(first.when_true)
            when_false = first.when_false
            for value in expression.values[1:]:
                next_condition = self._condition(value, when_false)
                when_true.extend(next_condition.when_true)
                when_false = next_condition.when_false
            return _Condition(entry_id, tuple(when_true), when_false)
        node_id = self._node("condition", _expression(expression), expression)
        self._connect(incoming, node_id)
        return _Condition(
            node_id,
            (_Tail(node_id, "when_true", "TRUE"),),
            (_Tail(node_id, "when_false", "FALSE"),),
        )

    def _converge(
        self,
        label: str,
        incoming: tuple[_Tail, ...],
        source: ast.AST,
    ) -> tuple[_Tail, ...]:
        if not incoming:
            return ()
        convergence_id = self._node("convergence", label, source)
        self._connect(incoming, convergence_id)
        return (_Tail(convergence_id),)

    def _node(self, kind: str, label: str, source: ast.AST) -> str:
        if len(self.nodes) >= self.max_nodes:
            raise DocForgeError(
                "logic_too_large",
                "A function exceeds the configured logic-node safety boundary",
                owner_node_id=self.owner_node_id,
                maximum=self.max_nodes,
            )
        line = max(1, int(getattr(source, "lineno", self.function.lineno)))
        column = max(0, int(getattr(source, "col_offset", 0)))
        self._sequence += 1
        logic_id = f"logic.{self._owner_digest}.{kind}.{line}.{column}.{self._sequence}"
        self.nodes.append(
            LogicNode(
                logic_id=logic_id,
                kind=kind,
                label=label,
                source_anchor=f"L{line}",
            )
        )
        return logic_id

    def _connect(self, incoming: tuple[_Tail, ...], target_id: str) -> None:
        for tail in incoming:
            self._edge(tail.source_id, tail.relation, target_id, tail.label)

    def _edge(
        self,
        source_id: str,
        relation: str,
        target_id: str,
        label: str | None,
    ) -> None:
        ordinal = self._edge_ordinals.get(source_id, 0)
        self._edge_ordinals[source_id] = ordinal + 1
        self.edges.append(
            LogicEdge(
                source_id=source_id,
                relation=relation,
                target_id=target_id,
                label=label,
                ordinal=ordinal,
            )
        )

    def _has_incoming(self, node_id: str) -> bool:
        return any(edge.target_id == node_id for edge in self.edges)


def _is_docstring(statement: ast.stmt) -> bool:
    return (
        isinstance(statement, ast.Expr)
        and isinstance(statement.value, ast.Constant)
        and isinstance(statement.value.value, str)
    )


def _is_catch_all(case: ast.match_case) -> bool:
    return (
        case.guard is None
        and isinstance(case.pattern, ast.MatchAs)
        and case.pattern.pattern is None
        and case.pattern.name is None
    )


def _expression(node: ast.AST) -> str:
    try:
        value = ast.unparse(node)
    except (AttributeError, ValueError):
        value = node.__class__.__name__
    return " ".join(value.split())


def _statement_label(statement: ast.stmt) -> str:
    if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef)):
        prefix = "async " if isinstance(statement, ast.AsyncFunctionDef) else ""
        return f"define {prefix}function {statement.name}"
    if isinstance(statement, ast.ClassDef):
        return f"define class {statement.name}"
    if isinstance(statement, ast.Pass):
        return "pass"
    return _expression(statement)


def _contains_runtime_call(statement: ast.stmt) -> bool:
    nested_definitions = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)
    stack: list[ast.AST] = [statement]
    while stack:
        current = stack.pop()
        if current is not statement and isinstance(current, nested_definitions):
            continue
        if isinstance(current, (ast.Call, ast.Await)):
            return True
        stack.extend(ast.iter_child_nodes(current))
    return False