Deconstructing Context Managers
"Managing context" is a fancy name for a simple concept. "Context" is a segment of code, and "managing" it means running something before and after it's executed. In the old days, it was done like this: ```python # Setup. try: # Do stuff. finally: # Teardown. ``` The setup and teardown are considered "managing", and the body of the try clause is considered "context". More often than not, the setup and teardown are related; like opening and closing a file, or acquiring and releasing a lock. But because of the way the language was structured, you couldn't encapsulate them in one place; The best you could do is put them in the same object, and invoke them before and after: ```python cm = ContextManager() cm.enter() try: # Do stuff. finally: cm.exit() ``` Then came the ``with`` statement, and instead of having to type all that you could just: ```python cm = ContextManager() with cm: # Do stuff. ``` And it would automatically call the context manager's ``__enter__`` method before and ``__exit__`` method after. Moreover, you could use the ``with ... as ...`` syntax to both start a managed context *and* bind a name to its context manager in a single breath: ```python cm = ContextManager() with cm as foo: # Do stuff. ``` The name is actually bound to ``__enter__``s return value, but it's usually the context manager object itself (i.e. ``__enter__`` usually returns ``self``). At first glance this looks stupid: why would we bind a new name for the context manager (``foo``) if we already have a name for it (``cm``)? But then you find out you can do this: ```python with ContextManager() as cm: # Do stuff. ``` That is, instantiate the context manager *in the ``with`` statement*, in which case its ``__init__`` is called, but its instance isn't bound to anything... Until it's returned by ``__enter__`` and bound to ``cm``! Of course, returning something other than ``self`` is also useful occasionally. Anyway, then the code happens, and then ``__exit__`` is called. If the context was exited cleanly, it receives three ``None``s, but if an exception was raised, it receives that exception's class, instance and traceback. In that case, if it returns ``True`` (or a Truthy value) that exception is suppressed, and if returns ``False`` (or a Falsey value like ``None``, which is the default) it continues to propagate. Let's see some examples! ![pause](/media/images/pause.png) Here's a context manager for a file: ```python class File(object): def __init__(self, path, mode='r'): self.path = path self.mode = mode self.file = None def __enter__(self): self.file = open(self.path, self.mode) return self def __exit__(self, exception, error, traceback): self.file.close() self.file = None # If we need the file once: with File(path) as f: # Do stuff. # And if need it multiple times: f = File(path) with f: # Do stuff. with f: # Do more stuff. ``` In fact, Python's ``file`` is already a context manager! ``` with open(path) as f: # Do stuff. ``` Except Python's ``file`` is opened *in its constructor*, not when the context starts (so its ``__enter__`` essentially does nothing), which means it can't be reused (because once it's closed, nothing reopens it). ![pause](/media/images/pause.png) Here's a context manager for a lock: ```python class Lock(object): def __init__(self): self.lock = threading.Lock() def __enter__(self): self.lock.acquire() return self def __exit__(self, exception, error, traceback): self.lock.release() # If we need the lock once (although it doesn't make much sense, # as locks are meant to be shared): with Lock(): # Do stuff. # And if we need it multiple times: lock = Lock() with lock: # Do stuff. with lock: # Do more stuff. ``` And again, Python's locks (``threading.Lock``, ``multiprocessing.Lock``, etc.) are already context managers. ```python lock = threading.Lock() with lock: # Do stuff. ``` ![pause](/media/images/pause.png) This is a cool one: a context manager that suppresses exceptions, so instead of doing this: ```python try: # Do stuff that may raise a ValueError. except ValueError: pass ``` We can do this: ```python with suppress(ValueError): # Do stuff that may raise a ValueError. ``` It's ridiculously simple: ```python class suppress(object): def __init__(self, *exceptions): self.exceptions = exceptions def __enter__(self): return self def __exit__(self, exception, error, traceback): return isinstance(error, self.exceptions) ``` And again, Python (at least Python ≥ 3.4) already has this in the standard ``contextlib`` module. Are you seeing the pattern yet? Context managers are the future. ```python with contextlib.suppress(ValueError): # Do stuff that may raise a ValueError. ``` It's actually quite an interesting module with a lot of fun ideas, so let's have a closer look at it. ![pause](/media/images/pause.png) The ``contextlib.closing`` context manager gets an object, and closes it once the context is exited, even if that object doesn't know to manage its own context. So while ``file`` does, what about ``socket``? ```python listener = socket.socket() try: listener.bind((host, port)) listener.listen(backlog) while True: connection, address = listener.accept() try: handle(connection) finally: connection.close() finally: listener.close() ``` Or... ```python listener = socket.socket() with contextlib.closing(listener): listener.bind((host, port)) listener.listen(backlot) while True: connection, address = listener.accept() with contextlib.closing(connection): handle(connection) ``` It turns out that since Python ≥ 3.2, ``socket`` was patched to manage its own context (along with a bunch of other stuff from the standard library, like ``urllib.open``), but the ``contextlib.closing`` utility still comes in handy sometimes. Here's how it would be implemented, since we're deconstructing stuff: ```python class closing(object): def __init__(self, target): self.target = target def __enter__(self): return self def __exit__(self, exception, error, traceback): self.target.close() ``` ![pause](/media/images/pause.png) The ``contextlib.redirect_stdout`` context manager replaces ``sys.stdout`` with a stream (``io.StringIO``), so if you have a function that prints something, you can test it: ```python def foo(): print('foo') stream = io.StringIO() with contextlib.redirect_stdout(stream): foo() assert stream.getvalue() == 'foo\n' ``` Or if you have some code littered with ``print``s, you can suppress them: ```python with contextlib.redirect_stdout(io.StringIO()): print('spam') print('spaaam') print('spaaaaam') ``` And in the spirit of deconstruction: ``` class redirect_stdout(object): def __init__(self, stream): self.stream = stream def __enter__(self): self.stdout = sys.stdout sys.stdout = self.stream return self def __exit__(self, exception, error, traceback): sys.stdout = self.stdout ``` ![pause](/media/images/pause.png) My absolute favorite, though, is the ``contextlib.contextmanager`` decorator. It's not a context manager in itself, but an amazing utility to make context management so much simpler - with generators! A generator is a function that runs in steps: ```python >>> def f(): ... print(1) ... yield ... print(2) ... yield >>> g = f() >>> next(g) 1 >>> next(g) 2 ``` And a context manager is a two-step process, right? So instead of having to write an entire class, with an (often trivial) ``__enter__`` and ``__exit__`` methods, we can do this: ```python @contextlib.contextmanager def context_manager(): # Setup yield # Teardown ``` And this will generate a class and everything for us. The value we yield is bound to the ``as`` name, and it even supports error handling: if an exception is raised, it's re-raised in the generator for us to handle. So a better template would be: ```python @contextlib.contextmanager def context_manager(): # Setup try: yield 'something' finally: # Teardown ``` This makes implementing almost all the context managers we've covered so far a breeze. Seriously, just have a look: ```python def file(path, mode='r'): f = open(path, mode) try: yield f finally: f.close() def suppress(*exceptions): try: yield except Exception as error: return isinstance(error, exceptions) def closing(target): try: yield finally: target.close() def redirect_stdout(stream): stdout = sys.stdout sys.stdout = stream try: yield finally: sys.stdout = stdout ``` The only problem is with context managers that are used multiple times, like locks, because ``contextlib.contextmanager`` requires the decorated generator to yield exactly once, so when we try to enter it again it's already "exhausted".But then again, most context managers are one time, so that's not that bad. Here's a deconstruction: ```python class contextmanager(object): def __init__(self, generator): self.generator = generator def __call__(self, *args, **kwds): self.execution = self.generator(*args, **kwds) def __enter__(self): return next(self.execution) def __exit__(self, exception, error, traceback): # No exception: if exception is None: next(self.execution) return # Inject the exception into the generator. try: self.execution.throw(exception, error, traceback) except Exception as e: # If the same exception propagated, keep it going. if e is error: return False # If a new exception occurred, re-raise it. raise else: # Otherwise, the exception was suppressed. return True ``` ![pause](/media/images/pause.png) This was the "lesson". In the [following post](/post/more-context-managers) we'll have some brand new, hopefully useful and probably wicked context managers.