Context manager able to write to the caller's namespace

[Chris Angelico]

On Fri, Jan 19, 2018 at 3:48 PM, Steven D'Aprano
<steve+comp.lang.python at pearwood.info> wrote:
> I want to define a context manager in one module:
>
> # a.py
> def CM:
>     def __enter__(self):
>         return self
>     def __exit__(self, *args):
>         pass
>
>
> Then call it from another module:
>
> # b.py
> import a
> with a.CM() as spam:
>     x = 1
>     y = 2
>
>
> in such a way that the spam context manager can, on exit, see the callers
> namespace and write to it. E.g. as a toy example (this isn't what I
> actually want to do!) we might have this:
>
> with a.CM() as spam:
>     x = 1
> print(x)
> # prints 2, not 1
>
> I stress that's not the intended functionality, it just demonstrates the
> requirement.

I'm assuming that you don't have to magically know that x was assigned
to, as that's its own separate problem. AIUI, you have three separate
cases:

1) Context manager was called from global scope, and needs access to
globals() or locals() as returned in the caller
2) Ditto ditto class scope, and needs access to the phantom
environment that's going to get put into the new class's dict
3) Ditto ditto function scope, which you specifically said is okay to
fail, but which also wants to be locals().

So basically, you want to call locals() in the caller's scope. I don't
think there's any standard way to do that, so you're going to be stuck
with sys._getframe() and the cross-implementation incompatibilities
that entails. But in CPython (tested in 3.7, should be fine in older
versions), this ought to work:

>>> def mutate():
...     locals = sys._getframe(1).f_locals
...     locals["x"] = 4
...
>>> x
1
>>> mutate()
>>> x
4
>>> class Foo:
...     x = 2
...     print(x)
...     mutate()
...     print(x)
...
2
4

Written as a context manager:

>>> import contextlib
>>> @contextlib.contextmanager
... def mutate():
...     yield
...     locals = sys._getframe(2).f_locals
...     locals["x"] = 4
...
>>> class Foo:
...     x = 2
...     print(x)
...     with mutate():
...         x = 3
...         print(x)
...     print(x)
...
2
3
4
>>>

There's basically zero guarantees about this though. Have fun. :)

ChrisA