Python: Standard Function And Context Manager?

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(

Solution 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:

classContextWrapper(object):def__init__(self, obj):
        self.__obj = obj

    def__enter__(self):
        returnselfdef__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:

classObjectWrapper(BaseClass):def__init__(self, obj):
        self.__class__ = type(
            obj.__class__.__name__,
            (self.__class__, obj.__class__),
            {},
        )
        self.__dict__ = obj.__dict__

    def__enter__(self):
        returnselfdef__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:

withclosing(my_func()) as my_obj:
    … do stuff …

Solution 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:

classmy_class(object):

    def__init__(self, result):
        print("this works", result)

classmanage_me(object):

    def__init__(self, callback):
        self.callback = callback

    def__enter__(self):
        return self

    def__exit__(self, ex_typ, ex_val, traceback):
        returnTruedef__call__(self, *args, **kwargs):
        return self.callback(*args, **kwargs)


defmy_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_medefmy_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)