Python - Fun With Context Managers

Context managers are one of those things in Python that can be very powerful and helpful. If you're used to C++, they can be roughly similar to a scoped_lock (or it's Boost predecessor). They are also exception-safe; your custom __exit__(), if provided, is guaranteed to be called to provide any possible cleanup. Some python programmers might not even know the term, but they understand that something "special" is happening when they write a with clause.

The most common use for context managers seems to be file I/O. That is often a casual programmer's first (and possibly last) introduction to them. It's great because you don't need to worry about closing your file explicitly:

with open('infile') as in_file:
    # work work work
# file is now closed

This is especially important if the next thing you are doing is launching a new process that will be reading from the file.

In larger applications, they're often used for properly thread locking (threading.Lock); thanks to C/C++ coders writing python, a code base I inherited had a lot of this:

# BAD EXAMPLE DO NOT USE
self._lock.acquire()
# work work work
self._lock.release()

Or worse:

# BAD EXAMPLE DO NOT USE
self._lock.acquire()
# work work work
if something:
    # complain
    self._lock.release()
    return
    
if something_else:
    # do something
    # two pages later... a call happens to fail...
    # forget to call self._lock.release()
    return

self._lock.release()

Sadly, that wasn't always in a try block, so you wouldn't always free up the lock if something exceptional happened. Nothing like random hangs to debug! Of course, the proper method, as noted in the official documentation, is:

with self._lock:
    # work work work
    # even an Exception can't stop the release

Writing your own context manager is as simple as creating an object that defines the dunder methods __enter__() and __exit__() with the proper signatures: __enter__(self) and __exit__(self, exc_type, exc_value, traceback). Like so many other things with Python, instead of dealing with the tedious, the standard library includes a wrapper that simplifies things: contextlib. The most straight-forward one is contextmanager which wraps a standalone function into a context manager with your function doing the "heavy lifting" and then yielding something: a piece of data, control of a shared resource, etc.

I highly recommend perusing the contextlib library reference page to see if some of the available wrappers could be useful in your day-to-day usage; for example contextlib.ExitStack makes rolling back much easier (and might be the subject of a future article here).

But I said "fun" in the title, so we'll play a little and use a context manager in a somewhat unique manner. A few months ago, I was helping with the packaging for an OSS project and refactored the python scripts they were using. I ended up using a context manager for console feedback:

from contextlib import contextmanager

class MsgPrinterException(Exception):
    """ Custom exception type """
    pass

GOOD_SYMBOL = 'Y' if platform.system() == "Windows" else '✓'
BAD_SYMBOL = 'N' if platform.system() == "Windows" else '✘'

@contextmanager
def msg_printer(start_msg=None):
    """ Context to use to print messages. Will exit if receives MsgPrinterException or re-raise any other """
    try:
        if start_msg:
            sys.stdout.write(f'> {start_msg} ')
            sys.stdout.flush()
            yield None
            sys.stdout.write('[' + GOOD_SYMBOL + ']\r\n')
            sys.stdout.flush()
        else:
            yield None
    except Exception as err:
        if start_msg:
            sys.stdout.write('[' + BAD_SYMBOL + ']\r\n')
        if isinstance(err, MsgPrinterException):
            print('\r\nERROR: ' + str(err))
            sys.exit(1)
        raise

Basically, if the code within the context doesn't blow up, it's interpreted as "good" and we output a check. If we throw our custom exception type, we skip the remaining code, say it's bad, output an X, and immediately exit the script. Any other exceptions, e.g. OSError, propagate to the caller.

with msg_printer('Setting up deployment directory...'):
    rm_path('deployment')
    os.makedirs('deployment')
    rm_path(f'bin/{target_arch}/Deploy')
    os.makedirs(f'bin/{target_arch}/Deploy')

def find_qt():
    """ Common code between Linux and Darwin """
    with msg_printer('Checking qtdir...'):
        if platform.system() == "Windows":
            raise MsgPrinterException("Windows shouldn't have called find_qt()!?!")

Thinking about using a context manager instead of try/except/else/finally may result in cleaner and easier-to-follow code, especially if it helps you avoid unnecessary duplication of similar statements. I think it's something that should be in your toolbox.

Aaron D. Marasco started programming by reading Turbo Pascal manuals in his front lawn during a mid '80s summer, has been using Unix and Linux since 1993 and 1999, and hates long walks on the beach. He has been a Jack-of-All-Trades and Engineering Catalyst at Innoplex since 2019.