Python `with` context vs generators/coroutines/tasks
I want to experiment with using python with blocks to apply modifiers to action within that block. But I'm not sure if it's possible to do this sensibly in the presence of coroutines.
For example, suppose I have a WithContext object that temporarily pushes onto a stack like this:
class WithContext:
stack = []
def __init__(self, val):
self.val = val
def __enter__(self):
WithContext.stack.append(self.val)
def __exit__(self, exc_type, exc_val, exc_tb):
WithContext.stack.pop()
def do_scoped_contextual_thing():
print(WithContext.stack[-1])
(Obviously the stack member would have to be thread-local, but ignore that for now.)
Then this code:
with WithContext("a"):
do_scoped_contextual_thing()
with WithContext("b"):
do_scoped_contextual_thing()
with WithContext("c"):
do_scoped_contextual_thing()
Will print:
a
b
c
But now suppose I have a coroutine situation:
def coroutine():
with WithContext("inside"):
yield 1
do_scoped_contextual_thing()
yield 2
with WithContext("outside"):
for e in coroutine():
do_scoped_contextual_thing()
print("got " + str(e))
I want this code to output:
outside
got 1
inside
outside
got 2
But actually it will output:
inside
got 1
inside
inside
got 2
The outsides changed to insides because the __enter__ inside the coroutine put a value on top of the stack, and __exit__ isn't called until the coroutine ends (instead of constantly enter-ing and exit-ing as you bounce in and out of the coroutine).
Is there a way to work around this problem? Are there "coroutine-local" variables?
I don't feel great about this, but I did modify your test code to re-enter the coroutine a few times. Similar to @CraigGidney's solution, this uses the inspect module to access and cache information on the call stack (aka, the "scope") in which a WithContext object is created.
I then basically search up the stack looking for a cached value, and use the id function to try and avoid holding references to the actual frame objects.
import inspect
class WithContext:
stack = []
frame_to_stack = {}
def __init__(self, val):
self.val = val
def __enter__(self):
stk = inspect.stack(context=3)
caller_id = id(stk[1].frame)
WithContext.frame_to_stack[caller_id] = len(WithContext.stack)
WithContext.stack.append( (caller_id, self.val))
def __exit__(self, exc_type, exc_val, exc_tb):
wc = WithContext.stack.pop()
del WithContext.frame_to_stack[wc[0]]
def do_scoped_contextual_thing():
stack = inspect.stack(context=0)
f2s = WithContext.frame_to_stack
for f in stack:
wcx = f2s.get(id(f.frame))
if wcx is not None:
break
else:
raise ValueError("No context object in scope.")
print(WithContext.stack[wcx][1])
def coroutine():
with WithContext("inside"):
for n in range(3):
yield 1
do_scoped_contextual_thing()
yield 2
with WithContext("outside"):
for e in coroutine():
do_scoped_contextual_thing()
print("got " + str(e))
One possible half-broken "solution" is to associate the context with a stack frame's location, and check for that location when looking up the context.
class WithContext:
_stacks = defaultdict(list)
def __init__(self, val):
self.val = val
def __enter__(self):
_, file, _, method, _, _ = inspect.stack()[1]
WithContext._stacks[(file, method)].append(self.val)
def __exit__(self, exc_type, exc_val, exc_tb):
_, file, _, method, _, _ = inspect.stack()[1]
WithContext._stacks[(file, method)].pop()
@staticmethod
def get_context():
for frame in inspect.stack()[1:]:
_, file, _, method, _, _ = frame
r = WithContext._stacks[(file, method)]
if r:
return r[-1]
raise ValueError("no context")
Note that constantly looking up stack frames is more expensive than just passing values around, and that you may not want to tell people you wrote this.
Note that this will still break in more complicated situations.
For example:
I had the same problem. Esentially I wanted to be able to execute code upon entering/leaving the running context of a coroutine, in my case to keep a call stack of sorts that is correct even in the case of interleaved yield s. It turns out tornado has support for this in the form of a StackContext which can be used as follows:
@gen.coroutine
def correct():
yield run_with_stack_context(StackContext(ctx), other_coroutine)
where ctx is a context manager that will enter and exit while the event loop is executing other_coroutine .
See https://github.com/muhrin/plumpy/blob/8d6cd97d8b521e42f124e77b08bb34c8375cd1b8/plumpy/processes.py#L467 for how I use it.
I've not looked into the implementation but tornado v5 switched to using asyncio as their default event loop so it should be compatible with that too.
上一篇: 状态机表示