Python: standard function and context manager?

The difficulty you’re going to run into is that for a function to be used as both a context manager (with foo() as x) and a regular function (x = foo()), the object returned from the function needs to have both __enter__ and __exit__ methods… and there isn’t a great way — in the general case — to add methods to an existing object.

One approach might be to create a wrapper class that uses __getattr__ to pass methods and attributes to the original object:

class ContextWrapper(object):
    def __init__(self, obj):
        self.__obj = obj

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        ... handle __exit__ ...

    def __getattr__(self, attr):
        return getattr(self.__obj, attr)

But this will cause subtle issues because it isn’t exactly the same as the object that was returned by the original function (ex, isinstance tests will fail, some builtins like iter(obj) won’t work as expected, etc).

You could also dynamically subclass the returned object as demonstrated here: https://stackoverflow.com/a/1445289/71522:

class ObjectWrapper(BaseClass):
    def __init__(self, obj):
        self.__class__ = type(
            obj.__class__.__name__,
            (self.__class__, obj.__class__),
            {},
        )
        self.__dict__ = obj.__dict__

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        ... handle __exit__ ...

But this approach has issues too (as noted in the linked post), and it’s a level of magic I personally wouldn’t be comfortable introducing without strong justification.

I generally prefer either adding explicit __enter__ and __exit__ methods, or using a helper like contextlib.closing:

with closing(my_func()) as my_obj:
    … do stuff …