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()

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.