How can I use a third-party context manager to run a custom-class's method?

Question

In python, I'd like to use a with-statement context manager when using a particular third-party package because their context manager will handle clean-up for me in a best-practice way.

Right now I have (say):

class MyClass(object):

    def __init__(self):
        self.open()

    def open(self):
        # Instantiate something *e.g.* open() a file

    def run(self):
        # Do some work *e.g.* read() from a filehandle

    def cleanup(self):
        # Clean everything up *e.g.* close() filehandle

This would then be called as

my_obj = MyClass()
my_obj.open()
my_obj.run()
my_obj.cleanup()

I'd ideally like to do all this using a context manager referenced within the class (if possible). It's the run() method that needs the context here. I am not sure how to do this - here are some theoretical attempts:

Option 1:

my_obj2=MyClass2()
with my_obj2.third_party_context_manager_method() as mcmm:
    my_obj2.run()

Option 2:

my_obj3=MyClass3()
my_obj3.run_modified_containing_third_party_context_manager()

Bonus: Now, in fact I need a double context manager in this case - how on earth..?

Please note: The file open example is just an example! My case is somewhat more complicated but the idea is that if I can make it work for that I can achieve anything :)

Another note: The use case is a dask cluster/client cleanup scenario - REF: How can I exit Dask cleanly?

Answer 1

I'm not sure if you have any limitations you have but here's a couple of basic ideas.

With a decorator:

def with_context(cm):
    def decorator(fn):
        def wrapper(*args, **kwargs):
            with cm:
                return fn(*args, **kwargs)
        return wrapper
    return decorator

class MyClass(object):
    ...
    @with_context(ThirdPartyCM)
    def run(self):
        # Do some work *e.g.* read() from a filehandle

Then just the simple way:

class MyClass(object):
    ...

    def _run(self):
        # Do some work *e.g.* read() from a filehandle

    def run(self):
        with ThirdPartyCM() as cm:
            self._run()

Alternatively you could inherit the 3rd party class if you needed it to be open for everything, but as you only need it open for run, I don't see why the above wouldn't work.