[Python] Can a context manager be used as a regular function?

Coding
#1

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.

#2

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
#3

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.

#4

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
#5

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.

#6

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.

#7

Haven’t tested it thoroughly but I think this is how I’d do it:


__author__ = 'Steve'

import os

class EnvSetting(object):
    def __init__(self, key, val):
        self._key = key
        self._val = str(val)
        self._cache = None

    def _restore(self, val):
        os.environ[self._key] = val

    def _delete(self):
        del os.environ[self._key]

    def set(self):
        if self._key in os.environ:
            current = os.environ.get(self._key)
            self._cache = lambda: self._restore(current)
        else:
            self._cache = lambda: self._delete()

        os.environ[self._key] = self._val

    def unset(self):
        self._cache()

    def __repr__(self):
        return "(os.environ['%s'] = '%s')" % (self._key, self._val)

class EnvCtx(object):
    def __init__(self, name, val):
        self.setter = EnvSetting(name, val)

    def __enter__(self):
        self.setter.set()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.setter.unset()

class EnvSetDecorator(object):
    def __init__(self, name, val):
        self.context = EnvCtx(name, val)

    def __call__(self, fn):
        def wrapped(*args, **kwargs):
            with self.context:
                return fn(*args, **kwargs)

        return wrapped

with EnvCtx('x', 123):
    print os.environ['x']


@EnvSetDecorator('x', 345)
def test():
    print os.environ['x']

test()

with EnvCtx('x', 999):
    print os.environ['x']
    test()

print os.environ.get('x', 'not found')
#8

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

#9

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' )
#10

Hard core! Though I think I’m not brave enough to go down that path myself :slight_smile:

#11

All I can say is, wow.

That is wonderfully crazy.