feat: try with exception edges, await and yield

README.md

@@ -64,8 +64,9 @@ PYTHONPATH=src python -m coretrace_python --emit-ir example.py
 Currently supported inside functions: parameters with defaults and keyword-only or star
 forms, decorators, assignments to names, attributes, items and unpacked tuples, augmented
 assignment, list, tuple and dict literals, `and`/`or`, chained comparisons, keyword
-arguments, `with`, `assert`, `if`/`elif`/`else`, `while`, `for`, `break`, `continue` and
-`raise`. Methods of module-level classes are analysed like functions; other module-level
+arguments, `with`, `assert`, `try`/`except`/`else`/`finally`, `await`, `yield`,
+`if`/`elif`/`else`, `while`, `for`, `break`, `continue` and `raise`. Blocks inside a `try`
+body carry exception edges to the handlers; `finally` is modelled on the normal path only. Methods of module-level classes are analysed like functions; other module-level
 code is skipped. A function using syntax outside this subset is reported by `--check` as
 an `unsupported-syntax` note and the other functions are still analysed. Each function is
 emitted as its control-flow graph: one block per basic block, ending in an explicit

src/coretrace_python/abstract/constants.py

@@ -113,6 +113,7 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat
     def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -> Mapping[BlockId, State]:
         block = self.blocks[block_id]
         state = self.evaluate(block, incoming)
+        exits = {target: state for target in block.exception_targets}
         terminator = block.terminator
         if isinstance(terminator, Branch):
             truth = state[terminator.condition].truthiness
@@ -121,12 +122,12 @@ def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -
                 Truth.FALSE: (terminator.else_block,),
                 Truth.UNKNOWN: (terminator.then_block, terminator.else_block),
             }[truth]
-            return {target: state for target in targets}
+            return {**exits, **{target: state for target in targets}}
         if isinstance(terminator, Jump):
-            return {terminator.target: state}
+            return {**exits, terminator.target: state}
         if isinstance(terminator, ForNext):
-            return {terminator.body: state, terminator.exit: state}
-        return {}
+            return {**exits, terminator.body: state, terminator.exit: state}
+        return exits
 
     # ------------------------------------------------------------------ transfer
 

src/coretrace_python/cfg/builder.py

@@ -35,6 +35,7 @@ def __init__(self, function: nodes.Function) -> None:
         self.blocks: dict[BlockId, BasicBlock] = {}
         self.counters: dict[str, int] = {}
         self.loops: list[tuple[BlockId, BlockId]] = []
+        self.handlers: list[tuple[BlockId, ...]] = []
 
     def build(self) -> CFG:
         entry = BlockId("entry")
@@ -50,10 +51,15 @@ def new_id(self, kind: str) -> BlockId:
         return BlockId(f"{kind}_{self.counters[kind]}")
 
     def finish(self, block: _Open, terminator: Terminator) -> None:
-        self.blocks[block.id] = BasicBlock(block.id, tuple(block.statements), terminator)
+        self.blocks[block.id] = BasicBlock(
+            block.id, tuple(block.statements), terminator, self.exception_targets()
+        )
 
     def header(self, block_id: BlockId, terminator: Terminator) -> None:
-        self.blocks[block_id] = BasicBlock(block_id, (), terminator)
+        self.blocks[block_id] = BasicBlock(block_id, (), terminator, self.exception_targets())
+
+    def exception_targets(self) -> tuple[BlockId, ...]:
+        return self.handlers[-1] if self.handlers else ()
 
     # ------------------------------------------------------------------ statements
 
@@ -104,9 +110,49 @@ def statement(
             return self.loop_statement(node, block, continuation)
         if isinstance(node, nodes.With):
             return self.with_statement(node, block)
+        if isinstance(node, nodes.Try):
+            return self.try_statement(node, block, continuation)
         block.statements.append(node)
         return block
 
+    def try_statement(
+        self, node: nodes.Try, block: _Open, continuation: BlockId | None
+    ) -> _Open | None:
+        """Body blocks carry exception edges to every handler; handlers, ``else`` and
+        ``finally`` join on the normal path (exceptional exits are approximated)."""
+
+        body_id = self.new_id("try")
+        handler_ids = tuple(self.new_id("handler") for _ in node.handlers)
+        else_id = self.new_id("else") if node.orelse else None
+        final_id = self.new_id("finally") if node.finalbody else None
+        after_id = continuation if continuation is not None and final_id is None else None
+        if after_id is None:
+            after_id = self.new_id("after") if final_id is None or continuation is None else None
+        join = final_id if final_id is not None else after_id
+        assert join is not None
+
+        self.finish(block, Jump(body_id, node.span))
+        self.handlers.append(handler_ids)
+        try:
+            # The whole body, including the block that leaves it, may raise into a handler.
+            end = self.sequence(node.body, _Open(body_id), None, node.span)
+            if end is not None:
+                self.finish(end, Jump(else_id if else_id is not None else join, node.span))
+        finally:
+            self.handlers.pop()
+        if end is not None and else_id is not None:
+            self.sequence(node.orelse, _Open(else_id), join, node.span)
+        for handler, handler_id in zip(node.handlers, handler_ids, strict=True):
+            opened = _Open(handler_id)
+            opened.statements.append(nodes.EnterHandler(handler, handler.span))
+            self.sequence(handler.body, opened, join, node.span)
+        if final_id is not None:
+            target = continuation if continuation is not None else after_id
+            assert target is not None
+            self.sequence(node.finalbody, _Open(final_id), target, node.span)
+            return None if continuation is not None else _Open(target)
+        return None if continuation is not None else _Open(after_id)  # type: ignore[arg-type]
+
     def with_statement(self, node: nodes.With, block: _Open) -> _Open | None:
         """Lay the body out inline; an early exit skips the ``ExitWith`` statements."""
 

src/coretrace_python/cfg/model.py

@@ -83,6 +83,8 @@ class BasicBlock:
     id: BlockId
     statements: tuple[nodes.Statement, ...]
     terminator: Terminator
+    exception_targets: tuple[BlockId, ...] = ()
+    """Handler blocks control may reach if a statement of this block raises."""
 
 
 @dataclass(frozen=True)
@@ -105,7 +107,9 @@ def __post_init__(self) -> None:
         for block_id, block in blocks.items():
             if block.id != block_id:
                 raise CFGError(f"block {block.id} is stored under {block_id}")
-            successors[block_id] = targets(block.terminator)
+            found = list(targets(block.terminator))
+            found.extend(t for t in block.exception_targets if t not in found)
+            successors[block_id] = tuple(found)
             for target in successors[block_id]:
                 if target not in blocks:
                     raise CFGError(f"block {block_id} targets unknown block {target}")

src/coretrace_python/frontend/ast_adapter.py

@@ -123,6 +123,13 @@ def expression(self, node: ast.expr) -> nodes.Expression:
                 for keyword in node.keywords
             )
             return nodes.Call(self.expression(node.func), arguments, keywords, span)
+        if isinstance(node, ast.Await):
+            return nodes.Await(self.expression(node.value), span)
+        if isinstance(node, ast.Yield):
+            yielded = self.expression(node.value) if node.value is not None else None
+            return nodes.Yield(yielded, span)
+        if isinstance(node, ast.YieldFrom):
+            self.fail(node, "yield from is not supported yet")
         if isinstance(node, (ast.ListComp, ast.SetComp, ast.GeneratorExp)):
             generators = tuple(self.generator(generator) for generator in node.generators)
             return nodes.Comprehension(
@@ -248,6 +255,24 @@ def statement(self, node: ast.stmt) -> nodes.Statement:
                 self.fail(node.cause, "raise from is not supported yet")
             exception = self.expression(node.exc) if node.exc is not None else None
             return nodes.Raise(exception, span)
+        if isinstance(node, ast.Try):
+            handlers = []
+            for handler in node.handlers:
+                handler_type = self.expression(handler.type) if handler.type is not None else None
+                handlers.append(
+                    nodes.ExceptHandler(
+                        handler_type, handler.name, self.block(handler.body), self.span(handler)
+                    )
+                )
+            return nodes.Try(
+                self.block(node.body),
+                tuple(handlers),
+                self.block(node.orelse),
+                self.block(node.finalbody),
+                span,
+            )
+        if isinstance(node, getattr(ast, "TryStar", ())):
+            self.fail(node, "except* groups are not supported yet")
         if isinstance(node, ast.Global):
             return nodes.Global(tuple(node.names), span)
         if isinstance(node, ast.Nonlocal):

src/coretrace_python/hir/nodes.py

@@ -103,6 +103,18 @@ class Dict:
     span: SourceSpan
 
 
+@dataclass(frozen=True, slots=True)
+class Await:
+    value: Expression
+    span: SourceSpan
+
+
+@dataclass(frozen=True, slots=True)
+class Yield:
+    value: Expression | None
+    span: SourceSpan
+
+
 @dataclass(frozen=True, slots=True)
 class ComprehensionGenerator:
     """One ``for target in iterable if condition...`` clause of a comprehension."""
@@ -136,6 +148,8 @@ class Comprehension:
     | Tuple
     | List
     | Dict
+    | Await
+    | Yield
     | Comprehension
 )
 
@@ -262,6 +276,33 @@ class ExitWith:
     span: SourceSpan
 
 
+@dataclass(frozen=True, slots=True)
+class ExceptHandler:
+    """``except type as name:``; ``type`` is ``None`` for a bare ``except``."""
+
+    type: Expression | None
+    name: str | None
+    body: tuple[Statement, ...]
+    span: SourceSpan
+
+
+@dataclass(frozen=True, slots=True)
+class Try:
+    body: tuple[Statement, ...]
+    handlers: tuple[ExceptHandler, ...]
+    orelse: tuple[Statement, ...]
+    finalbody: tuple[Statement, ...]
+    span: SourceSpan
+
+
+@dataclass(frozen=True, slots=True)
+class EnterHandler:
+    """Synthetic statement the CFG builder emits at the start of a handler block."""
+
+    handler: ExceptHandler
+    span: SourceSpan
+
+
 @dataclass(frozen=True, slots=True)
 class ImportAlias:
     name: str
@@ -330,6 +371,8 @@ class Class:
     | With
     | EnterWith
     | ExitWith
+    | Try
+    | EnterHandler
     | Import
     | ImportFrom
     | Global

src/coretrace_python/interprocedural/modulegraph.py

@@ -99,10 +99,13 @@ def _prefixes(dotted: str) -> list[str]:
 def _statements(body: Iterable[nodes.Statement]) -> Iterable[nodes.Statement]:
     for statement in body:
         yield statement
-        for attribute in ("body", "orelse"):
+        for attribute in ("body", "orelse", "finalbody"):
             nested = getattr(statement, attribute, None)
             if isinstance(nested, tuple):
                 yield from _statements(nested)
+        if isinstance(statement, nodes.Try):
+            for handler in statement.handlers:
+                yield from _statements(handler.body)
 
 __all__ = [
     "IGNORED_DIRECTORIES",

src/coretrace_python/interprocedural/summaries.py

@@ -185,14 +185,15 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat
     def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -> Mapping[BlockId, State]:
         block = self.blocks[block_id]
         state = self.evaluate(block, incoming)
+        exits = {target: state for target in block.exception_targets}
         terminator = block.terminator
         if isinstance(terminator, Branch):
-            return {terminator.then_block: state, terminator.else_block: state}
+            return {**exits, terminator.then_block: state, terminator.else_block: state}
         if isinstance(terminator, Jump):
-            return {terminator.target: state}
+            return {**exits, terminator.target: state}
         if isinstance(terminator, ForNext):
-            return {terminator.body: state, terminator.exit: state}
-        return {}
+            return {**exits, terminator.body: state, terminator.exit: state}
+        return exits
 
     # ------------------------------------------------------------------ transfer
 

src/coretrace_python/ir/lowering.py

@@ -18,6 +18,7 @@
 from coretrace_python.hir.visitors import Node, children
 from coretrace_python.ir.model import (
     Assert,
+    Await,
     BasicBlock,
     BinaryOp,
     BoolOp,
@@ -26,6 +27,7 @@
     BuildList,
     BuildTuple,
     Call,
+    Catch,
     Compare,
     Constant,
     EffectInstruction,
@@ -51,6 +53,7 @@
     ValueInstruction,
     WithEnter,
     WithExit,
+    Yield,
 )
 from coretrace_python.semantic import SEMANTIC_ANALYSES
 from coretrace_python.semantic.scopes import (
@@ -157,6 +160,11 @@ def expression(self, node: nodes.Expression) -> Value:
         if isinstance(node, nodes.Dict):
             items = tuple((self.expression(k), self.expression(v)) for k, v in node.items)
             return self.emit(BuildDict(self.new_value(), node.span, items))
+        if isinstance(node, nodes.Await):
+            return self.emit(Await(self.new_value(), node.span, self.expression(node.value)))
+        if isinstance(node, nodes.Yield):
+            value = self.expression(node.value) if node.value is not None else None
+            return self.emit(Yield(self.new_value(), node.span, value))
         if isinstance(node, nodes.Attribute):
             object_value = self.expression(node.value)
             return self.emit(GetAttr(self.new_value(), node.span, object_value, node.name))
@@ -212,6 +220,13 @@ def statement(self, node: nodes.Statement) -> None:
         if isinstance(node, nodes.ExitWith):
             self.emit_effect(WithExit(None, node.item.span, self.contexts[node.item.span]))
             return
+        if isinstance(node, nodes.EnterHandler):
+            handler = node.handler
+            caught_type = self.expression(handler.type) if handler.type is not None else None
+            caught = self.emit(Catch(self.new_value(), handler.span, caught_type))
+            if handler.name is not None:
+                self.store(nodes.Name(handler.name, handler.span), caught)
+            return
         if isinstance(node, nodes.ExpressionStatement):
             self.expression(node.expression)
             return
@@ -276,7 +291,11 @@ def function(self, node: nodes.Function) -> FunctionIR:
             for statement in cfg_block.statements:
                 self.statement(statement)
             terminator = self.terminator(cfg_block)
-            blocks.append(BasicBlock(cfg_block.id, tuple(self.instructions), terminator))
+            blocks.append(
+                BasicBlock(
+                    cfg_block.id, tuple(self.instructions), terminator, cfg_block.exception_targets
+                )
+            )
         return FunctionIR(
             qualified_name(self.scopes, node), parameter_values, self.cfg.entry, tuple(blocks), node.span
         )
@@ -309,6 +328,8 @@ def names(target: nodes.Target | None) -> None:
     def walk(node: Node) -> None:
         if isinstance(node, nodes.Assign | nodes.AugAssign | nodes.For | nodes.WithItem):
             names(node.target)
+        elif isinstance(node, nodes.ExceptHandler) and node.name is not None:
+            assigned.add(node.name)
         if isinstance(node, nodes.Function | nodes.Class | nodes.Comprehension):
             return
         for child in children(node):