Multiple ways to invoke context manager in python

Question

---

Background

I have a class in python that takes in a list of mutexes. It then sorts that list, and uses __enter__() and __exit__() to lock/unlock all of the mutexes in a specific order to prevent deadlocks.

The class currently saves us a lot of hassle with potential deadlocks, as we can just invoke it in an RAII style, i.e.:

self.lock = SuperLock(list_of_locks)
# Lock all mutexes.
with self.lock:
    # Issue calls to all hardware protected by these locks.

---

Problem

We'd like to expose ways for this class to provide an RAII-style API so we can lock only half of the mutexes at once, when called in a certain way, i.e.:

self.lock = SuperLock(list_of_locks)
# Lock all mutexes.
with self.lock:
    # Issue calls to all hardware protected by these locks.

# Lock the first half of the mutexes in SuperLock.list_of_locks
with self.lock.first_half_only:
    # Issue calls to all hardware protected by these locks.

# Lock the second half of the mutexes in SuperLock.list_of_locks
with self.lock.second_half_only:
    # Issue calls to all hardware protected by these locks.

---

Question

Is there a way to provide this type of functionality so I could invoke with self.lock.first_half_only or with self.lock_first_half_only() to provide a simple API to users? We'd like to keep all this functionality in a single class.

Thank you.

Answer 1

Yes, you can get this interface. The object that will be entered/exited in context of a with statement is the resolved attribute. So you can go ahead and define context managers as attributes of your context manager:

from contextlib import ExitStack  # pip install contextlib2
from contextlib import contextmanager

@contextmanager
def lock(name):
    print("entering lock {}".format(name))
    yield
    print("exiting lock {}".format(name))

@contextmanager
def many(contexts):
    with ExitStack() as stack:
        for cm in contexts:
            stack.enter_context(cm)
        yield

class SuperLock(object):

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

    def __enter__(self):
        # implement for entering the `with self.lock:` use case
        return self

    def __exit__(self, exce_type, exc_value, traceback):
        pass

    @property
    def first_half_only(self):
        return many(self.list_of_locks[:4])

    @property
    def second_half_only(self):
        # yo dawg, we herd you like with-statements
        return many(self.list_of_locks[4:])

When you create and return a new context manager, you may use state from the instance (i.e. self).

Example usage:

>>> list_of_locks = [lock(i) for i in range(8)] 
>>> super_lock = SuperLock(list_of_locks) 
>>> with super_lock.first_half_only: 
...     print('indented') 
...   
entering lock 0
entering lock 1
entering lock 2
entering lock 3
indented
exiting lock 3
exiting lock 2
exiting lock 1
exiting lock 0

---

Edit: class based equivalent of the lock generator context manager shown above

class lock(object):

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

    def __enter__(self):
        print("entering lock {}".format(self.name))
        return self

    def __exit__(self, exce_type, exc_value, traceback):
        print("exiting lock {}".format(self.name))
        # If you want to handle the exception (if any), you may use the
        # return value of this method to suppress re-raising error on exit