Python: standard function and context manager?

Question

In python, there are many functions that work as both standard functions and context managers. For example open() can be called either as:

my_file=open(filename,'w')

or

with open(filename,'w') as my_file:

Both giving you a my_file object that can be used to do whatever you need to. In general the later is preferable, but there are times when one might want to do the former as well.

I've been able to figure out how to write a context manager, either by creating a class with __enter__ and __exit__ functions or by using the @contextlib.contextmanager decorator on a function and yield rather than return. However, when I do this I can no longer use the function straight - using the decorator, for example, I get a _GeneratorContextManager object back rather than the result that wanted. Of course, if I made it as a class, I'd just get an instance of the generator class, which I'd assume is essentially the same thing.

So how can I design a function (or class) that works as either a function, returning an object, or a context manager, returning a _GeneratorContextManager or the like?

edit:

For example, say I have a function like the following (this is HIGHLY simplified):

def my_func(arg_1,arg_2):
    result=arg_1+arg_2
    return my_class(result)

So the function takes a number of arguments, does stuff with them, and uses the result of that stuff to initialize a class, which it then returns. End result is I have an instance of my_class, just like I would have a file object if I had called open. If I want to be able to use this function as a context manager, I can modify it like so:

@contextlib.contextmanager
def my_func(arg_1,arg_2):
    result=arg_1+arg_2 # This is roughly equivalent to the __enter__ function
    yield my_class(result)
    <do some other stuff here> # This is roughly equivalent to the __exit__function

Which works just fine when calling as a context manager, but I no longer get an instance of my_class when calling as a straight function. Perhaps I'm just doing something wrong?

Edit 2:

Note that I do have full control over my_class, including the ability to add functions to it. From the accepted answer below, I was able to infer that my difficulty stemmed from a basic misunderstanding: I was thinking that whatever I called (my_func in the example above) needed to have the __exit__ and __enter__ functions. This is not correct. In fact, it's only what the function returns (my_class in the above example) that needs the functions in order to work as a context manager.

Answer 1

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 …

Answer 2

Just for clarity: if you are able to change my_class, you would of course add the __enter__/__exit__ descriptors to that class.

If you are not able to change my_class (which I inferred from your question), this is the solution I was referring to:

class my_class(object):

    def __init__(self, result):
        print("this works", result)

class manage_me(object):

    def __init__(self, callback):
        self.callback = callback

    def __enter__(self):
        return self

    def __exit__(self, ex_typ, ex_val, traceback):
        return True

    def __call__(self, *args, **kwargs):
        return self.callback(*args, **kwargs)


def my_func(arg_1,arg_2):
    result=arg_1+arg_2
    return my_class(result)


my_func_object = manage_me(my_func) 

my_func_object(1, 1)
with my_func_object as mf:
    mf(1, 2)

---

As a decorator:

@manage_me
def my_decorated_func(arg_1, arg_2):
    result = arg_1 + arg_2
    return my_class(result)

my_decorated_func(1, 3)
with my_decorated_func as mf:
    mf(1, 4)