fix: match correct try block markers in CFG (#205)

* fix: match correct try block markers in CFG

Fix an ``AssertionError`` in stack-size computation when a single exception
region is split into multiple ``TryBegin`` instances sharing one handler (as
produced by bytecode-rewriting tools that wrap a whole function body in a
single handler). ``TryBegin``/``TryEnd`` are now matched on the handler block
rather than on ``TryBegin`` identity.

* Update doc/changelog.rst

Co-authored-by: Matthieu Dartiailh <marul@laposte.net>

* fix version in changelog

---------

Co-authored-by: Matthieu Dartiailh <marul@laposte.net>

Author: P403n1x87

doc/changelog.rst

@@ -1,6 +1,17 @@
 ChangeLog
 =========
 
+03-06-2026: Version 0.18.1
+--------------------------
+
+Bugfixes:
+
+- Fix an ``AssertionError`` in stack-size computation when a single exception
+  region is split into multiple ``TryBegin`` instances sharing one handler (as
+  produced by bytecode-rewriting tools that wrap a whole function body in a
+  single handler). ``TryBegin``/``TryEnd`` are now matched on the handler block
+  rather than on ``TryBegin`` identity.
+
 03-06-2026: Version 0.18.0
 --------------------------
 

src/bytecode/cfg.py

@@ -343,7 +343,17 @@ def run(self) -> Generator[Union[_StackSizeComputer, int], int, None]:
                 # current try begin. However inside the CFG some blocks may
                 # start with a TryEnd relevant only when reaching this block
                 # through a particular jump. So we are lenient here.
-                if instr.entry is not self._current_try_begin:
+                #
+                # We match on the exception handler (the TryBegin target block)
+                # rather than on the TryBegin instance: a single exception region
+                # can be split into several TryBegin copies that share the same
+                # handler (see ``from_bytecode``), and the copy carried over as
+                # ``pending_try_begin`` through a jump is not necessarily the same
+                # instance as the one referenced by this block's leading TryEnd.
+                if (
+                    self._current_try_begin is None
+                    or instr.entry.target is not self._current_try_begin.target
+                ):
                     continue
 
                 # Compute the stack usage of the exception handler
@@ -403,7 +413,7 @@ def run(self) -> Generator[Union[_StackSizeComputer, int], int, None]:
                     if (
                         (te := self.block.get_trailing_try_end(i))
                         and self._current_try_begin is not None
-                        and te.entry is self._current_try_begin
+                        and te.entry.target is self._current_try_begin.target
                     ):
                         assert isinstance(te.entry.target, BasicBlock)
                         yield from self._compute_exception_handler_stack_usage(

tests/test_cfg.py

@@ -18,6 +18,7 @@
     Label,
     SetLineno,
     TryBegin,
+    TryEnd,
     dump_bytecode,
 )
 from bytecode.utils import PY312, PY313, PY314
@@ -903,6 +904,78 @@ def test(*args, **kwargs):
             self.assertEqual(test([], name=None), -1)
             self.assertEqual(stdout.getvalue(), "second finally\nfirst finally\n")
 
+    def test_stack_size_computation_shared_handler_split_regions(self):
+        # Regression test for an assertion failure in the stack-size
+        # computation. A bytecode-rewriting transform (e.g. an instrumentation
+        # tool) may wrap a whole function body in a single exception handler
+        # whose coverage is split around the body's own try blocks, so that the
+        # regions never nest but all share the *same* handler block. A single
+        # logical region then maps to several TryBegin instances.
+        #
+        # When such a region is exited through a forward jump (here the branch
+        # inside the first handler body), the jump target block is reached both
+        # by fall-through and by the jump, and ends up carrying one TryBegin
+        # copy as ``pending_try_begin`` while its own leading TryEnd references
+        # a *different* copy of the same region. Matching TryEnd against the
+        # current TryBegin by identity then failed to close the region, leaving
+        # it spuriously open and tripping ``assert self._current_try_begin is
+        # None`` at the following TryBegin. Matching must be done on the handler
+        # (the TryBegin target block) instead.
+        if sys.version_info < (3, 11):
+            self.skipTest("exception tables (TryBegin/TryEnd) require 3.11+")
+
+        def trigger(d, x):  # pragma: no cover
+            try:
+                d["a"]
+            except KeyError:
+                if x:  # a branch *inside* the handler body
+                    pass
+                else:
+                    d["b"] = 1
+            try:  # ... followed by another try block
+                del d["c"]
+            except KeyError:
+                pass
+            return 99
+
+        bc = Bytecode.from_code(trigger.__code__)
+
+        # Wrap the whole body in one shared handler, splitting the covering
+        # region around the body's existing try blocks.
+        except_label = Label()
+        first_tb = last_tb = TryBegin(except_label, push_lasti=True)
+        i = 0
+        while i < len(bc):
+            instr = bc[i]
+            if isinstance(instr, TryBegin) and last_tb is not None:
+                bc.insert(i, TryEnd(last_tb))
+                last_tb = None
+                i += 1
+            elif isinstance(instr, TryEnd):
+                j = i + 1
+                while j < len(bc) and not isinstance(bc[j], TryBegin):
+                    if isinstance(bc[j], Instr):
+                        last_tb = TryBegin(except_label, push_lasti=True)
+                        bc.insert(i + 1, last_tb)
+                        break
+                    j += 1
+                i += 1
+            i += 1
+        bc.insert(0, first_tb)
+        bc.append(TryEnd(last_tb))
+        bc.append(except_label)
+        bc.append(Instr("PUSH_EXC_INFO"))
+        bc.append(Instr("RERAISE", 2))
+
+        cfg = ControlFlowGraph.from_bytecode(bc)
+        # Used to raise AssertionError on 3.12.
+        cfg.compute_stacksize()
+
+        # The rewritten code object is valid and still runs correctly.
+        trigger.__code__ = cfg.to_code()
+        self.assertEqual(trigger({}, True), 99)
+        self.assertEqual(trigger({"c": 1}, False), 99)
+
     def test_stack_size_with_dead_code(self):
         # Simply demonstrate more directly the previously mentioned issue.
         def test(*args):  # pragma: no cover