-
-
Save JelleZijlstra/a53b17417c5189b487316628acc5555f to your computer and use it in GitHub Desktop.
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() |
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 thewith
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.
For reference:
This gets displayed by
dis
as: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.