"""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 = "", 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