I have a function that sets some values and I sometimes want to use it in a with-statement (so it restores the original values after I do some stuff). I couldn’t find a way to do this aside from making the contextmanager a class, but I don’t know if that is pythonic? The behavior I am thinking of would mimic the open() function:
import contextlib
@contextlib.contextmanager
def f():
print 'do some stuff'
yield
print 'exiting the context!'
f()
# do some stuff
with f():
pass
# do some stuff
# exiting the context!
I’m assuming the correct approach is to make a separate function and a context manager, but I figured I would check first.
So if you look in the contextlib.py file, you’ll see that the contextmanager decorator actually uses a class behind the scenes.
Now if you want an object that can both act as a context manager, and a function, just define a call method in addition to the enter and exit methods.
class ContextTest(object):
def __init__(self, *args):
self.args = args
def __enter__(self):
print('Entering a context block.')
return self(*self.args)
def __call__(self, *args):
print('Calling: {0}'.format(args))
return self
def __exit__(self, type, value, traceback):
print('Exiting the context block')
return False
# So this the following
with ContextTest('Hello') as context:
context('World')
# Should result in:
# Entering a context block.
# Calling: ('Hello',)
# Calling: ('World',)
# Exiting the context block
# Vs something like this
context = ContextTest()
with context('Hello') as c:
pass
# Results in:
# Calling: ('Hello',)
# Entering a context block.
# Calling: ()
# Exiting the context block
That’s not quite what I’m looking for. I want a function that can be called outside of a with statement and do some stuff, and can be called in a with statement and do some stuff and then undo that stuff. I think I’m trying to over design a function and I’m just going to have two separate methods: one that is a context manager, and one that isn’t.
So that would just be a regular context manager then right?
Basically if you are preforming your work in init, that will always happen whether its a with block or not.
import os
class SetEnv(object):
def __init__(self, key, value):
self._key = key
os.environ[key] = value
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
os.environ.pop(self._key)
return False
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# No key named SomeKey
with SetEnv('SomeKey', 'SomeValue'):
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# SomeValue
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# No key named SomeKey
SetEnv('SomeKey', 'SomeValue')
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# SomeValue
Now granted you don’t have to use a class for this, you can technically get away with the decorator method, you’d just need a different wrapper for the non-context manager version.
from contextlib import contextmanager
# Please don't ever actually do something like this.
# Its super hacky
def set_env(key, value):
os.environ[key] = value
yield
os.environ.pop(key)
# using the decorator directly instead of the special syntax
SetEnv = contextmanager(set_env)
def notcontext(func):
# Only calling next once, will mean anything below the yield line won't execute
return func.next()
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# No key named SomeKey
with SetEnv('SomeKey', 'SomeValue'):
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# SomeValue
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# No key named SomeKey
notcontext(set_env('SomeKey', 'SomeValue'))
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# SomeValue
I think the premise is not wise – a function should be a function and a context manager a context manager. Relying on side-effects always leads to trouble down the road: two maintainers who see that object in two different contexts will assume it does two very different things. If you really want to keep them together:
class Context():
....
@staticmethod
def functional_part():
....
with Context():
Context.functional_part()
which is one more line but a lot clearer. If there is some kind of state you could make functionalpart() an instance method instead:
class Context():
....
def functional_part(self):
....
with Context() as ctx:
ctx.functional_part()
but in that case it’s more like ‘context-manager-with-optional-extra-function’ because it’s an antipattern to make a Context() you don’t want just to get it’s instancemethod.
I do agree, the making a context manager, just to get a throw-away instance is really ugly.
I do like the ability to have a manager that will do auto-cleanup, but if you use the class without it, being able to cleanup after yourself manually, open being a good example of this.
So I really should have setup the context manager more like:
import os
class SetEnv(object):
def __init__(self, key, value):
self._key = key
os.environ[key] = value
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.cleanup()
return False
def cleanup(self):
os.environ.pop(self._key)
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# No key named SomeKey
with SetEnv('SomeKey', 'SomeValue'):
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# SomeValue
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# No key named SomeKey
changed_env = SetEnv('SomeKey', 'SomeValue')
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# SomeValue
changed_env.cleanup()
print(os.environ.get('SomeKey', 'No key named SomeKey'))
# No key named SomeKey
I can’t figure out a clean way to do have a similar system off of the contextmanager decorator though, I get some really weird behavior by calling next manually multiple times. It seems to silently stop the execution flow somehow.
Pymel and mGui’s layout classes are good examples of adding optional functionality with context managers. You can handle the parent stack using the context managers, or if you just want to create a new layout and manually parent it, you can.
Yeah, I like that. Breaks everything up into small logical chunks.
I also really like having both the context manager, and a decorator built from that manager.
I’ve got a few of those for some of the annoying Maya quirks (I’m looking at you auto-keyframe).
For academic purpose only, to explore some possibilities…
from opcode import HAVE_ARGUMENT
from inspect import currentframe
from dis import opmap
from contextlib import contextmanager
def callable_contextmanager( func ):
def wrapper( arg ):
caller = currentframe().f_back
last_opcode = caller.f_code.co_code[caller.f_lasti]
# Warning! The following might not cover all possible cases
last_size = 3 if ord( last_opcode ) > HAVE_ARGUMENT else 1
calling_opcode = caller.f_code.co_code[caller.f_lasti + last_size]
if ord( calling_opcode ) == opmap["SETUP_WITH"]:
# Called within a context manager
return contextmanager( func )( arg )
else:
# Called as a regular function
return func( arg ).next()
return wrapper
Usage
# from some_module import callable_contextmanager
@callable_contextmanager
def f( arg ):
print 'do some stuff ' + arg
yield
print 'exiting the context!'
with f( 'as context' ) as stuff:
pass
yielded = f( 'as regular function' )
