PEP 789

pep789.py

from contextlib import contextmanager
import dis
import sys

from contextlib import nullcontext

@contextmanager
def block_yields():
    calling_frame = sys._getframe(2)
    offset = calling_frame.f_lasti
    bc = dis.Bytecode(calling_frame.f_code)
    # This is the chunk of the exception table that corresponds to the block
    # of code that starts the with block.
    exc_entry = None
    # TODO this only works if there are no nested with blocks
    # The right solution probably involves repeatedly following
    # the target of the exception entry, and checking all blocks that
    # eventually point to the .target instruction of the exc_entry.
    for e in bc.exception_entries:
        if e.start == offset + 2:
            exc_entry = e
            break
    assert exc_entry is not None, "No exception entry found"
    for instr in bc:
        if instr.offset < exc_entry.start:
            continue
        if instr.offset >= exc_entry.end:
            break
        if instr.opname == "YIELD_VALUE":
            raise RuntimeError(f"Cannot yield from within a block (line {instr.line_number})")
    yield


@contextmanager
def assert_raises(exc_type):
    try:
        yield
    except exc_type:
        pass
    else:
        raise AssertionError(f"Expected {exc_type} to be raised")


def fail1():
    with block_yields():
        yield 42

def pass1():
    yield 1
    with block_yields():
        print("no yield")
    yield 2

def pass2():
    yield 1
    with block_yields():
        with nullcontext():
            pass
    yield 2

def fail2():
    with block_yields():
        with nullcontext():
            yield 42


def test():
    dis.dis(fail2)
    with assert_raises(RuntimeError):
        list(fail1())
    list(pass1())
    list(pass2())
    # TODO
    # with assert_raises(RuntimeError):
    #     list(fail2())


if __name__ == "__main__":
    test()

For reference:

• exception table entries look like this:

>>> dis.Bytecode(block_yields.pass2)
Bytecode(<function pass2 at 0x103b26a50>)
>>> _.exception_entries
[_ExceptionTableEntry(start=4, end=34, target=170, depth=0, lasti=True), _ExceptionTableEntry(start=34, end=56, target=136, depth=1, lasti=True), _ExceptionTableEntry(start=56, end=58, target=102, depth=2, lasti=True), _ExceptionTableEntry(start=60, end=76, target=136, depth=1, lasti=True), _ExceptionTableEntry(start=76, end=102, target=170, depth=0, lasti=True), _ExceptionTableEntry(start=102, end=122, target=130, depth=4, lasti=True), _ExceptionTableEntry(start=122, end=136, target=136, depth=1, lasti=True), _ExceptionTableEntry(start=136, end=156, target=164, depth=3, lasti=True), _ExceptionTableEntry(start=156, end=170, target=170, depth=0, lasti=True)]

This gets displayed by dis as:

ExceptionTable:
  L1 to L2 -> L16 [0] lasti
  L2 to L3 -> L12 [1] lasti
  L3 to L4 -> L8 [2] lasti
  L5 to L6 -> L12 [1] lasti
  L6 to L8 -> L16 [0] lasti
  L8 to L10 -> L11 [4] lasti
  L10 to L12 -> L12 [1] lasti
  L12 to L14 -> L15 [3] lasti
  L14 to L16 -> L16 [0] lasti

I think this format was new in 3.11, so for earlier version a different approach is necessary.

I tried first to simply iterate over the bytecode, increment a counter on BEFORE_WITH, and decrement it on WITH_EXCEPT_START, but that doesn't work (at least on main) because some of the bytecode gets reordered. Possibly this approach will work before 3.11 though.

Thanks Jelle! I've mentioned this in python/peps#3782, albeit as a 'rejected alternative'. I'm confident that it could support syntatically-nested with-statements, but:

it's not yet clear how this would work when user-defined context managers wrap sys.prevent_yields. Worse, this approach ignores explicit calls to __enter__() and __exit__(), meaning that the context management protocol would vary depending on whether the with statement was used.

The 'only pay if you use it' performance cost is very attractive. However, inspecting frame objects is prohibitively expensive for core control-flow constructs, and causes whole-program slowdowns via de-optimization. On the other hand, adding interpreter support for better performance leads back to the same pay-regardless semantics as our preferred solution above.

Good point about direct __enter__ calls. I think this approach may still be useful for earlier versions of Python, but it's definitely a hack that falls down in some contexts.

Absolutely - I was thinking about shipping it in Trio, albeit only active in debug mode for perf reasons. I think we're pretty likely to do this eventually, at least for python 3.11 - 3.13 where we have the exceptions table but no native solution.