Writing/The Unbearable Lightness of Coding
Chapter 14 · 36 min read
Objects 5: Lifecycle and Context Management
In this chapter, we'll talk about the object lifecycle: how it's created and deleted, and why this process is somewhat unreliable. We'll see how these limitations can be overcome with context managers, which let us control an object's life to an unprecedented resolution — covering the nuances of the syntax, different ways this abstraction can improve our software design, and various tricks that we can use to write better context managers.
***
Let's start with the object lifecycle. We've discussed objects at length and in depth, but we never really addressed how they are created or destroyed — in part because it's mostly managed by Python, behind the scene.
However, there are actually some magic methods that represent this behavior, too: we've already seen __init__ multiple times. Still, this method, which we compared to the constructor, doesn't actually construct anything — after all, it receives the object, already constructed, via the self argument:
>>> class A:
... def __init__(self):
... print('init', self)
>>> a = A()
init <A object at 0x...> # object already exists
Which is why I also called it an initializer — a method that, as it name implies, merely setting up the object's initial state.
Creating Objects
For an actual constructor, we have __new__, which is called before the object exists to actually create it. Since there's no object to put in self yet, __new__ is a class method, even if we don't decorate it with @classmethod: its first argument is going to be the class that's being instantiated:
>>> class A:
... def __new__(cls):
... print('new', cls)
>>> a = A()
new <class 'A'>
Note that if we define both of them, for some reason only __new__ happens:
>>> class A:
... def __new__(cls):
... print('new', cls)
... def __init__(self):
... print('init', self)
>>> a = A()
new <class 'A'>
But that's because in our current implementation, __new__ doesn't return anything — and so, implicitly, returns None, on which Python refrains from invoking __init__. To return the object, however, we have to somehow create it; and since we're implementing its very creation, we're a bit stuck. This recursive deadend is similar to the one we encountered when trying to access an attribute from within __getattribute__ — which we solved, if you recall, by admitting we can't do it ourselves, and delegating to Python by invoking super (which, if we don't have any explicit superclasses, object will handle). Like I said, __new__ is a classmethod — but because it's undecorated (and anyway, predates classmethod's class-binding), we have to pass it in ourselves:
>>> class A:
... def __new__(cls):
... print('new', cls)
... self = super().__new__(cls)
... print('self', self)
... return self
... def __init__(self):
... print('init', self)
>>> a = A()
new <class 'A'>
self <A object at 0x...>
init <A object at 0x...>
Also note that __new__ and __init__ have to have the same signature; if we do:
>>> class A:
... def __new__(cls):
... return super().__new__(cls)
... def __init__(self, x):
... self.x = x
>>> a = A(1)
TypeError: A.__new__() takes 1 positional argument but 2 were given
We naturally get an error, since both these functions are effectively called on the same occasion and with the same arguments — namely, an instantiation of a class, like A(1):
>>> class A:
... def __new__(cls, x):
... return super().__new__(cls)
... def __init__(self, x):
... self.x = x
>>> a = A(1)
And while here, our constructor doesn't do anything, we can in principle initialize the object there:
>>> class A:
... def __new__(cls, x):
... self = super().__new__(cls)
... self.x = x
... return self
... def __init__(self, x):
... print('init', self.x)
>>> a = A(1)
init 1
Thus rendering __init__ redundant; we just need to remember to return the result at the end, which would be considered the newly created object.
Pooling Objects
What's the point of having __new__, though, if it's the same as __init__, only more cumbersome? In most cases, there is no point — which is exactly why __init__ exists; it's that initialization bit of an already existing object that we're generally interested in.
However, if we want to intervene in the creation process itself — for example, if we want to reuse objects that we've already created — __new__ is really our only option. Imagine something like a class representing a file, where we pass in a path and get back an object with all sorts of useful functionality. If we'd want to improve its efficiency or consistency, we might notice that each instance represents a file on disk — so really, having multiple instances that point to the same file is a waste, and runs the risk of them falling out of sync with each other. So instead:
>>> class File:
... pool = {}
... def __new__(cls, path):
... path = os.path.realpath(path)
... if path in cls.pool:
... return cls.pool[path]
... self = super().__new__(cls)
... cls.pool[path] = self
... return self
... def __init__(self, path):
... self.path = path
Here, we keep a pool of all our objects, and in __new__, we intercept the path, normalize it (so we can tell if different relative paths actually resolve to the same thing), and check if it's there. If not, we go ahead with creating the object — but before we return it, we also add it to the pool; so the next time a similar path is passed in, we can fetch the existing object and reuse it. Here it is in action:
>>> f1 = File('file.txt')
>>> f1
<File object at 0x000002E9528DB7C0>
>>> f2 = File('file.txt')
>>> f2
<File object at 0x000002E9528DB7C0>
>>> f1 is f2
True
Pooling Issues
So far, so good — but there's a problem: because of the otherwise-convenient __new__ and __init__ separation, we end up reinitializing objects that were already created, which is usually not what we want. Suppose we had an attribute of whether the file is open, which we initially set to False:
>>> class File:
... # same as before
... def __init__(self, path):
... # same as before
... self.is_open = False
At some point (when the file is opened), it would change to True:
>>> f1 = File('file.txt')
>>> f1.is_open = True
However, the moment a similar file is created — although __new__ is smart enough to reuse this object — __init__ gets called again, overwriting this attribute and accidentally resetting the object's status:
>>> f2 = File('file.txt')
>>> f1.is_open
False
This happens since really, Python has no way of telling why we're overriding __new__: whether we're returning new objects, old ones, or something else altogether. In principle we could do something like:
>>> class A:
... def __new__(cls):
... return 1
>>> a = A()
>>> a
1
So the only thing it keeps track of is whether the result is an instance of A, and if so, invokes its __init__. If we're serious about having a pool of objects, then, we should account for that — like, by moving all our initialization code to __new__:
>>> class File:
... pool = {}
... def __new__(cls, path):
... path = os.path.realpath(path)
... if path in cls.pool:
... return cls.pool[path]
... self = super().__new__(cls)
... self.path = path
... self.is_open = False
... cls.pool[path] = self
... return self
... # no __init__
Which prevents reinitialization, on account of the second __new__ returning the existing object without having reaching this part. Still, we should be cautious with subclasses: those can't have __init__s now, either, and even if they implement their own __new__:
>>> class JSONFile(File):
... def __new__(cls, path):
... self = super().__new__(cls, path)
... self.is_parsed = False
... return self
Their reinitialization would still happen, because while their superclass skips resetting the is_open attribute for existing object, they have no way of knowing whether is_parsed should be set (or reset) or not. A more reliable solution might be settings a special flag for cached objects upon their first initialization, and bailing if it's set:
>>> class File:
... pool = {}
... def __new__(cls, path):
... path = os.path.realpath(path)
... if path in cls.pool:
... return cls.pool[path]
... self = super().__new__(cls)
... self._init = False
... cls.pool[path] = self
... return self
... def __init__(self, path):
... if self._init:
... return
... self.path = path
... self.is_open = False
... self._init = True
It's not very elegant, but at least this clause can be replicated in subclasses:
>>> class JSONFile(File):
... def __init__(self, path):
... if self._init:
... return
... super().__init__(path)
... self.is_parsed = False
Deleting Objects
In any case, now we know the first to steps in the object's lifecycle: it's created in __new__ and initialized in __init__. After that, it goes on to live its life, storing some attributes, being passed around methods and so on. And eventually, it gets deleted — at which point its resources are released and its removed from memory. For this last occasion we have the __del__ magic method, known as the object's destructor. If we store temporary files we need to remove, or connections we need close, it might be tempting to think of it as the right place to do so — but actually, it's not ideal.
You see, because Python manages object allocation and reclamation for us, it's not guaranteed when an object is actually deleted; in other words, __del__ can be less predictable than we'd like. Not that its logic is undefined; it's just rather complicated. This, for example, works as expected:
>>> class A:
... def __del__(self):
... print('deleting', self)
>>> a = A()
>>> del a
deleting <A object at 0x...>
But this doesn't:
>>> a = A()
>>> b = a
>>> del a
# nothing
Because the object is still referenced by b, and therefore still alive. Once it isn't:
>>> del b
deleting <A object at 0x...>
Then the destructor is called. Python has a whole component, called the garbage collector, in charge of recognizing, terminating and reclaiming unnecessary objects — and it turns out to be a pretty sophisticated process, with many optimizations and nuances in place.
We're not going to get into its details, but very roughly speaking, it uses reference counting: it keeps track of what objects reference what other objects, and only ever removes something from memory if it is no longer referenced by anyone else. In other words, only once an object becomes unreachable does Python agree to release it, knowing it's safe to say nobody else is going to use (or even reach) it.
Weak References
In simple cases, like with our a and b, these references can be easily traced; but in other cases it's not so obvious. Remember our File class from before? Its pool is a dictionary that keeps references to all its instances, so they can't be deleted, ever — their reference count never drops to zero, so the garbage collector never reclaims their memory and __del__ is never invoked:
>>> class File:
... # same as before
... def __del__(self):
... print('deleting', self)
>>> f = File('file.txt')
>>> del f
# nothing
To be fair, Python has no way of telling whether we actually intended for this dictionary to keep all these instances alive; if we didn't, we should have used the weakref module of the standard library. It lets us wrap an object in a weak reference, which provides a point to it without increasing its reference count, so it doesn't interfere with garbage collection. Here's how it works:
>>> import weakref
>>> a = A()
>>> b = weakref.ref(a)
>>> b()
<A object at 0x...>
But now, if we delete a, at some point the garbage collector will notice it and reclaim its memory — in which case b will still be there, but calling it will return None. Notice that this doesn't happen immediately:
>>> del a
# nothing
>>> b()
<A object at 0x...>
Because the garbage collector runs periodically in the background, not wanting to interfere with our program's flow and degrade its performance. If we want to invoke it directly, we can do so with the standard gc module:
>>> import gc
>>> gc.collect()
deleting <A object at 0x...>
>>> b()
None
If we decide our file pool shouldn't keep file objects alive, then, we could do:
>>> class File:
... # same as before
... def __new__(cls, path):
... path = os.path.realpath(path)
... if path in cls.pool:
... ref = cls.pool[path]
... obj = ref()
... if obj is None:
... print('need to recreate object')
... else:
... return obj
... self = super().__new__(cls)
... cls.pool[path] = weakref.ref(self)
... self.path = path
... self.is_open = False
... return self
In which case we get:
>>> f = File('file.txt')
>>> del f
deleting <File object at 0x...>
The weak reference is still in the pool, mind you, but it's a hollow shell; and we'll notice it as soon as we instantiate a similar path, which will prompt creation rather than reuse:
>>> f = File('file.txt')
need to recreate object
Admittedly, it's a bit clunky; and since caching dictionaries are such a common use case, weakref provides as with WeakKeyDictionary and WeakValueDictionary that implement the same logic for us. In our case, its the values of the dictionary whose references we want to weaken, so all we needed to change in our original implementation was:
>>> class File:
... pool = weakref.WeakValueDictionary()
... # same as before
And we'd get the same behavior:
>>> f = File('file.txt')
>>> del f
deleting <File object at 0x...>
>>> dict(File.pool)
{}
Better, even, since you can see that the dictionary knows to clean itself up.
Circular References
Even weak references are not always enough, though. Suppose we have a class that represents a directory, keeping a list of all its files, with each file pointing back to its directory via a parent attribute. None of these references are weak, per say; and since they are circular, their reference count will never drop to zero, as they effectively keep each other alive. I'll demonstrate with a simpler example:
>>> class C:
... def __init__(self):
... self.d = D(self)
... def __del__(self):
... print('deleting', self)
>>> class D:
... def __init__(self, c):
... self.c = c
... def __del__(self):
... print('deleting', self)
>>> c = C()
>>> del c
# nothing
c points to d, which points back to c, so garbage collection is not trivial. We could add a method that would unlink the two manually, and call it right before deleting:
>>> class C:
... # same as before
... def unlink():
... self.d.unlink()
... self.d = None
>>> class D:
... # same as before
... def unlink(self):
... self.c = None
>>> c = C()
>>> c.unlink()
deleting <D object at 0x...>
>>> del c
deleting <C object at 0x...>
But we need to remember calling it on any occasion where a might go out of scope — which means we're essentially reintroducing memory management to a language that is supposed to take care of it for us.
Alternatively, we might decide that circular references are an anti-pattern and vow never to use them — but you know where I'm going with this: in some cases, they make a whole lot of sense, and result in much more powerful and intuitive abstractions. Besides, even if we avoid them entirely, Python itself has plenty. A function, for example, has a __globals__ dictionary that might very well contain it:
>>> def f():
... pass
>>> f.__globals__['f']
<function f at 0x...>
And that's a circular reference right there. It may be a harmless one, since del f means deleting f from precisely this globals dictionary — but there are trickier cases. An raised error, for example, retains a trackeback pointing to all the frames it propagated, which in turn point to all their local variables. If this error is somehow retained:
>>> def f():
... a = A()
... raise ValueError()
>>> try:
... f()
... except ValueError as error:
... saved_error = error
# nothing
Then as you can see, a's deletion message is missing, since it's still alive deep down in:
>>> saved_error.__traceback__.tb_next.tb_frame.f_locals['a']
<A object at 0x...>
Digging Through Garbage
So yeah, anyone who's vehemently against circular references should probably use a different language altogether. In Python's defense, it goes out of its way to mitigate all these edge cases; errors, for example, normally go out of scope at the end of their except clause, so their traceback (and all the locals, in turn) do get collected; as you see, I had to capture it in a separate name to preserve it. Even circular references can be collected: the garbage collection periodically goes over all the objects, trying to detect circles that are the only thing keeping themselves alive. Our previous c and d, for example, would be reclaimed eventually, even without unlinking — we just needed to prompt the garbage collector if we wanted to see it happen immediately:
>>> c = C()
>>> del c
>>> gc.collect()
deleting <C object at 0x...>
deleting <D object at 0x...>
I like the garbage collector: it's pretty sophisticated machinery that constantly operates in the background, dealing with the fundamental question of memory management on our behalf — and yet, most of the time, it's so well designed that we don't even know it's there. Even when we do, it's delightfully transparent, exposing us to its inner workings with honest, straightforward tools. Take get_referrers, for example: it returns a list of all the objects referencing its argument:
>>> c = C()
>>> gc.get_referrers(c.d)
[{'d': <__main__.D object at 0x...>}]
Which can be a bit surprising, since we'd expect to see c in this list; where did this odd dictionary come from? But upon reflection, we might recognize it after all:
>>> c.__dict__
[{'d': <__main__.D object at 0x...>}]
That is, c actually doesn't point to d — at least, not directly; it points to its __dict__, which then points to d. If we ask for that __dict__'s referrers, we can confirm our theory:
>>> gc.get_referrers(gc.get_referrers(c.d))
[<C object at 0x...>]
Simiarly, we can get all the referents of an object — that is, objects that it refers to:
>>> gc.get_referents(c)
[{'d': <__main__.D object at 0x...>}, <class '__main__.C'>]
And we see that really, an object only points to its __dict__ and __class__ — everything else, like we discovered with __getattribute__, can be inferred from them.
There's a lot more to tell about garbage collection, but this book is pretty long as it is. One last gem I'd like to leave you with, though, is this:
>>> x = None
>>> class Lazarus:
... def __del__(self):
... global x
... x = self
... print('deleting', x)
This is called a Lazarus object, because when it dies, it effectively gets resurrected — as it's being deleted, it binds itself to a new name:
>>> y = Lazarus()
>>> del y
deleting <Lazarus object at 0x…>
>>> x
<Lazarus object at 0x...>
And so invariably stays alive. This is undoubtedly what its developer wanted, for whatever reason — so Python delivers; the problems start, however, when it gets deleted again. In a more naive implementation, this would invoke __del__ again, which would revive the object again, and worse — if the destructor closes or frees some resources, such a repeated cleanup may very well crash. So, Python introduced the idea of a finalized object — one that still exist, but whose destructor was already called:
>>> y = Lazarus()
>>> gc.is_finalized(y)
False
>>> del y
deleting <Lazarus object at 0x…>
>>> x
<Lazarus object at 0x...>
>>> gc.is_finalized(x)
True
>>> del x
# nothing happens (__del__ doesn't repeat)
This way, while the object does survive its deletion (mostly likely to be further inspected or logged), it also gets marked — so that on its next deletion, it's gone for good. I've never seen anything like it in any production code, but I find the mere fact that Python even addresses this case (let alone, resolves it elegantly) is a testament to just how thoughtful and thorough its design is.
Setup and Teardown
Overall, Python's garbage collection is great, but it's not perfect: we can't always predict the order in which destructors are called, and some circular references might elude it. The conclusion, though, is not to manage memory ourselves, nor avoid circular references altogether; it's only __del__ that we should use cautiously, without relying on its exact timing.
This is pretty bad news. Sometimes, we have objects that require some setup, or prologue, before they become operational, and some teardown, or epilogue, after we're done with them. A file, for example, has to be opened before we read it — and afterwards, it needs to be close. Doing this:
>>> def read(path):
... f = open(path)
... print(f.read())
... f.close()
Works well in the happy flow, when there are no problems — but if read fails (e.g. due to a sudden disk failure), it might raise an exception that will cut our function short, never closing the file. It might feel insignificant, but it's still a memory leak; and if we don't account for our resources, letting too many of them slip between the cracks without being released, their dead weight can add up — clogging our memory over time, slowing our system, maybe even crashing it. The proper way to read a file would be:
>>> def read(path):
... f = open(path)
... try:
... print(f.read())
... finally:
... f.close()
That is, try reading from it, and finally — no matter what happens, whether it works or an exception is raised — close it, reclaiming its memory.
A more substantial example might be thread synchronization. We'll talk about concurrent programming at length in a future chapter: for now, suffice to say that sometimes, we want different parts of the program to happen in parallel, in so called threads; and sometimes, those threads need to be synchronized. Just like people in a talking circle wouldn't want to talk over each other, threads shouldn't write over the same variable at the exact same time — so we isolate their access to shared data in what's known as a critical section, and use a special object called a lock to protect it. In a talking circle, it's akin to a symbolic token: only the person holding some stick or pillow is allowed to speak, until they pass it on to somebody else. Similarly, each thread should acquire the lock before entering its critical section, which is guaranteed to happen instantly for exactly one thread — and after it's done, it should release the lock, letting other threads that might be waiting for it to get their shot:
>>> import threading
>>> lock = threading.Lock()
>>> def thread():
... lock.acquire()
... # critical section
... lock.release()
In this case, however, if the thread ends or raises an error without releasing the lock, it's not a small matter of inefficient memory usage; the entire program can grind to a halt — like losing the symbolic token, so everyone in the circle are just sitting in silence indefinitely. The proper way to write handle a critical section, then, is like so:
>>> import threading
>>> lock = threading.Lock()
>>> def thread():
... lock.acquire()
... try:
... # critical section
... finally:
... lock.release()
Now, in other languages, there's a more natural way to handle such logics: setup happens in the constructor, and teardown happens in the destructor. For a file, it's pretty straightforward — we open it at the beginning and close it at the end, so we don't need to worry ourselves with trys and finallys; the moment it goes out of scope, it's resources are released. For something like a lock, there's a nice trick known as resource acquisition is initalization, or RAII: we can wrap the lock in an auxiliary object that acquires it in its constructor and releases it in its destructor, thus guaranteeing that this will happen the moment the function ends one way or another:
>>> class LockGuard:
... def __init__(self, lock):
... self.lock = lock
... self.lock.acquire()
... def __del__(self):
... self.lock.release()
>>> def thread():
... guard = LockGuard(lock)
... # critical section
In Python, however, we saw that we don't have the benefit of reliable destruction — so we can't count on when exactly the file will be closed, or the lock released. It seems that Python's automatic memory management, despite its best efforts, fundamentally compromised our ability to control the object lifecycle, and thus to have proper setup and teardown logics; it seems we are doomed to repeat our prologues and epilogues in clumsy verbosity, cramming our code between ungainly trys and finallys.
Context Management
Unless, that is, we find a different way to think about this altogether. If we're honest, we don't really care about the object's lifecycle — it's not a question of life and death, and it doesn't matter when the object gets created or deleted. All we want is for this object to go through some phase: to manage what happens before and after it gets into a particular context. Construction and destruction are not essential to this task: they're just natural moments to hook such logic into, since one happens before anything, and the other after everything.
So instead, Python provides us with separate, dedicated hooks to do just that, and an exceedingly elegant syntax to invoke them. Any object that implements them is called a context manager, and can be placed in a with-statement:
>>> def thread():
... with lock:
... # critical section
Before the nested code block begins, the context manager is prompted to set it up; and after it ends, whether successfully or with an error, the context manager is called upon to tear it down and decide how to proceed. Naturally, locks implement this interface to be acquired and released, as do file objects:
>>> def read(path):
... with open(path) as f:
... print(f.read())
Here, the file is actually constructed by the builtin open function, as part of the with-statement, so an as-clause is used to bind it to a name.
And of course, we can write context managers ourselves! The setup part happens in __enter__, and the teardown in __exit__, whose signature we'll cover in a minute — for now, just to showcase how simpe it is:
>>> class A:
... def __enter__(self):
... print('setting up', self)
... def __exit__(self, *args):
... print('tearing down', self)
>>> a = A()
>>> with a:
... print('in context', a)
setting up <A object at 0x...>
in context <A object at 0x...>
tearing down <A object at 0x...>
The Entering Nuance
If we try to do this, however:
>>> with A() as a:
... print('in context', a)
setting up <A object at 0x...>
in context None
tearing down <A object at 0x...>
We find out that the as-clause doesn't actually bind to the result of instantiation — that's the entering nuance. When this clause is present, it captures the return value of __enter__: so to support object creation directly in a with-statement, it's customary to return the object itself:
>>> class A:
... def __enter__(self):
... print('setting up', self)
... return self
... # same as before
>>> with A() as a:
... print('in context', a)
setting up <A object at 0x...>
in context <A object at 0x...>
tearing down <A object at 0x...>
This is not necessary, however; it's easy enough to conceive of a context manager for something like a temporary directory, which creates one upon entry, deletes it afterwards, and yields not itself, but the created path:
>>> import os, shutil, tempfile
>>> class TemporaryDirectory:
... def __enter__(self):
... self.path = tempfile.mkdtemp()
... return self.path
... def __exit__(self, *args):
... if os.path.exists(self.path):
... shutil.rmtree(self.path)
>>> with TemporaryDirectory() as tempdir:
... print(os.path.exists(tempdir))
True
>>> os.path.exists(tempdir)
False
The Exit Signature
Exiting has a different nuance, for which we need to discuss its signature, first: it receives the error triplet. If the context ended successfully, with no error, all three values are None — but if it was cut short by a raised exception, the first value would be that exception's class, the second would be the exception itself, and the third would be a traceback object:
>>> class A:
... def __enter__(self):
... return self
... def __exit__(self, exception, error, traceback):
... print(exception, error, traceback)
>>> with A():
... pass
None None None
>>> with A():
... raise ValueError('something went wrong')
<class 'ValueError'> something went wrong <traceback object at 0x...>
ValueError: something went wrong
There is something a bit redundant about it being three separate values: the error class is available in error.__class__, and its traceback in error.__traceback__; but that's how Python developed historically, so that's what we have.
In some cases, whether the context ended successfully doesn't matter — a file should be closed and a lock should be released regardless; and as you can see, the exception propagates, so enclosing scopes can catch and handle it.
But when the context's success does matter, here we have a chance to inspect it and act accordingly. Consider a database, for example: sometimes, a series of operations need to happen together, or not at all — like when transferring money from one bank account to another. The trivial solution would be to subtract the sum from the sender's account and add it to the receiver's; but if the system crashes right in the middle, we might end up with the money being lost altogether (or created out of thin air, if we switch the order). For that reason, most database provide us with the ability to start a transaction: a series of operations that definitely happen together when it's committed, or don't happen at all if it's rolled back. In code, it would look like this:
>>> db = Database()
>>> tx = db.begin_transaction()
>>> try:
... tx.update(record1).set(balance=record1.balance - sum)
... tx.update(record2).set(balance=record2.balance + sum)
... tx.commit()
... except Exception:
... tx.rollback()
But we can implement it much nicer if we turn the transaction object returned from begin_transaction into a conext manager:
>>> class Transaction:
... # implementation
... def __enter__(self):
... return self
... def __exit__(self, exception, error, traceback):
... if error:
... self.rollback()
... else:
... self.commit()
>>> with db.begin_transaction() as tx:
... tx.update(account1).set(balance=account1.balance - sum)
... tx.update(account2).set(balance=account2.balance + sum)
The Exit Nuance
Another interesting feature of __exit__ is that besides inspecting the error, it can also suppress it. If its return value (which is None by default, like any function) is false, the exception is propagated — but if it returns True, the error is suppressed, and the code after it continues uninterrupted. This effectively lets us have a local return-statement, similar to goto END, skipping the rest of our code block:
>>> class ContextExit(Exception):
... pass
>>> class GotoEnd:
... def __enter__(self):
... return self
... def __exit__(self, exception, error, traceback):
... return isinstance(error, ContextExit)
... def end(self):
.... raise ContextExit()
So one more way to break out of a nested loop would be:
>>> with GotoEnd() as goto:
... for x in [1, 2, 3]
... for y in [3, 4, 5]
... print(x, y)
... if x == y:
... goto.end()
1 3
1 4
1 5
2 3
2 4
2 5
3 3
This gives rise to some other interesting use cases, too. For example, you know how sometimes you want some code to run, but don't care if it raises an exception?
>>> try:
... # code
... except Exception:
... pass
Or, we saw that running an infinite loop, like a server, is often interrupted with control+C:
>>> try:
... run_server()
... except KeyboardInterrupt:
... pass
We can support both cases more elegantly if we create a context manager that suppresses particular exceptions (or all of them, if none are specified):
>>> class Suppress:
... def __enter__(self, *exceptions):
... self.exceptions = exceptions or Exception
... def __enter__(self):
... return self
... def __exit__(self, exception, error, traceback):
... return isinstance(error, self.exceptions)
As an aside, isinstance (and issubclass — and some other builtin functions in Python, like a string's startswith) accept either an argument or a tuple thereof, returning true if the object is an instance of (or a subclass of, or starts with) any of the provided options. Now, we can do this:
>>> with Suppress():
... # code
>>> with Suppress(KeyboardInterrupt):
... run_server()
Which is pretty neat, isn't it?
Context Prevalence
In fact, we don't have to implement it ourselves — a similar context manager is already part of the standard library as functools.suppress. And just like files and locks can manage their contexts, so can all the standard sockets, iterables, and most other examples we've seen. In fact, the standard way to create temporary directories is precisely in a with-statement, and using a database (like SQLite3) implicitly creates transactions that are committed on success and rolled back on error:
>>> with tempfile.TemporaryDirectory() as tempdir:
... print(os.path.exists(tempdir))
True
>>> os.path.exists(tempdir)
False
>>> import sqlite3
>>> db = sqlite3.connect('db.sqlite3')
>>> with db: # happens in a transaction
... db.execute('UPDATE accounts SET balance = balance - ? WHERE id = ?',
... (sum, account_id1))
... db.execute('UPDATE accounts SET balance = balance + ? WHERE id = ?',
... (sum, account_id2))
Generally, context prevalence were a big hit. The idea might have originated in other languages, like Lisp's with-open-file or C#'s using statement, but I think it's fair to say Python really perfected and popularized the concept — and it's yet another example of how open-minded its design is, assimilating the best of all worlds and consolidating it under its flexible data model.
So while our discussion started with the somewhat unpredictable object lifecycle, we ended up with an abstraction that's much more potent, and can be used in a surprising variety of situations. One scenario I like in particular is that of testing code, so I'd like to give a couple more examples from this domain. Normally, testing in Python is done with assert-statement:
>>> def div(x, y):
... return x / y
>>> def test_div():
... assert div(4, 2) == 2
... assert div(1, 2) == 0.5
If the expression after the assert is true, it passes silently, and if not raises an AssertionError. But how do we test exceptions, like division by zero? Doing this:
>>> def test_div_error():
... try:
... div(1, 0)
... assert False
... except ZeroDivisionError:
... assert True
Works, but it's pretty horrendous. Luckily, with our newfound power, we can have:
>>> class AssertRaises:
... def __init__(self, exception):
... self.exception = exception
... def __enter__(self):
... return self
... def __exit__(self, exception, error, traceback):
... assert error is not None, 'no exception raised'
... return isinstance(error, self.exception)
Which would complain if no exception was raised, proceed happily if the one we were expecting was, and propagate any other error — all with two beautiful lines:
>>> def test_div_error():
... with AssertRaises(ZeroDivisionError):
... div(1, 0)
Another common challenge in testing code is dealing with standard output. If a function returns a string, it's easy enough to test its contents are correct — but if it goes ahead and prints it out, how can we capture it? Sometimes, we can tackle this ahead of time with software design; instead of:
>>> def greet(name):
... print('hello', name or 'stranger')
We can separate the string creation from its emission:
>>> def create_greeting(name):
... return 'hello ' + (name or 'stranger')
>>> def greet(name):
... print(greeting(name))
Which makes testing greet redundant, since we can safely assume print works as expected, and focus on the easier-to-test create_greeting instead. But sometimes, it's not so easy to decouple — so instead, we need to patch Python's sys.stdout, the file-like object into which print writes. We can replace it with another file-like object, io.StringIO, which stores anything written to it internally, making it possible for us to fetch it later — and just as importantly, after the test is done, put the original stream back in place (or it might swallow our test report altogether):
>>> import io, sys
>>> class ReplaceStdout:
... def __init__(self, exception):
... self.stream = io.StreaingIO()
... def __enter__(self):
... self.stdout = sys.stdout
... sys.stdout = self.stream
... return self.stream
... def __exit__(self, exception, error, traceback):
... sys.stdout = self.stdout
And now, behold:
>>> def test_greet():
... with ReplaceStdout() as stream:
... greet('alice')
... assert stream.getvalue() == 'hello alice\n'
Definitely not we had in mind when we just started talking about construction and destruction — and yet, it works marvelously.
Context Decorators
As you might've noticed already, there's a certain affinity between context managers and decorators: both are ways to inject custom logic before and after something happens — whether a function invocation or a code block. In one case, we have:
>>> def decorator(f):
... def wrapper(*args, **kwargs):
... print('before')
... try:
... return f(*args, **kwargs)
... finally:
... print('after')
... return wrapper
>>> @decorator
... def f():
... print('inside')
>>> f()
before
inside
after
And on the other:
>>> class ContextManager:
... def __enter__(self):
... print('before')
... def __exit__(self, exception, error, traceback):
... print('after')
>>> def f():
... with ContextManager():
... print('inside')
>>> f()
before
inside
after
Both options have their appeal: decorators can be added with an elegant @ on top and leverage all sorts of functional trickery, while context managers provide a cleaner separation between before and after, and can be used on any code block, a resolution finer than function calls. Luckily, we don't have to choose: the standard contextlib module provides a ContextDecorator superclass which can turn our context manager into a decorator for free:
>>> import contextlib
>>> class ContextManager(contextlib.ContextDecorator):
... # same as before
>>> @ContextManager()
>>> def f():
... print('inside')
>>> f()
before
inside
after
And, knowing everything we know about objects, it's quite trivial to imagine a possible implementation:
>>> class ContextDecorator:
... def __call__(self, f):
... def wrapper(*args, **kwargs):
... self.__enter__()
... try:
... result = f(*args, **kwargs)
... self.__exit__(None, None, None)
... return result
... except Exception as error:
... exc, tb = error.__class__, error.__traceback__
... if not self.__exit__(exc, error, tb):
... raise # re-raises the same exception
... return wrapper
Context Generators
This is nice, but contextlib has something much cooler in store. Turning a context manager object into a decorator function is straightforward, since we have __call__ to emulate invocations; but what about the other way around? Can we turn a function into a context manager?
Initially, it seems impossible: a function has a single behavior (__call__), whereas the context manager provides two: setup (__enter__) and teardown (__exit__). At the same time, these behaviors happen before and after the managed context; and if you recall the structure we started with:
# setup
try:
# context
finally:
# teardown
This block sure fits nicely within a function. And although it was cumbersome to spell out every time, we replaced it with an elegant with-statement already. As far as implementing a context manager goes, however, this is somewhat cleaner than defining a class with two magic methods — if only there was a way to write a function that stops in the middle, right before the context, and can resume later right after it…
Enter contextlib.contextmanager: a mad decorator that lets us turn a generator with exactly one yield-statement into a context manager class. If this is the first time you see this, it might look bizarre, but bear with me — it's quite genius. Our previous temporary directory example could be actually implemented like this:
>>> @contextlib.contextmanager
... def temporary_directory():
... path = tempfile.mkdtemp()
... try:
... yield path
... finally:
... if os.path.exists(path):
... shutil.rmtree(path)
That is, we create a temporary directory like before, then yield its path — that's what gets assigned in the as-clause, same as the return value of __enter__ — and finally clean up, also like before. As for our database transaction example:
>>> class Database:
... # implementation
... @contextlib.contextmanager
... def begin_transaction():
... tx = Transaction()
... try:
... yield
... tx.commit()
... except Exception:
... tx.rollback()
... raise
Or, if we reimplement Suppress:
>>> @contextlib.contextmanager
... def suppress(*exceptions):
... try:
... yield
... except exceptions or Exception:
... pass
Except-clauses, it turns out, also accept a tuple of options. In any case, like you can see, it's a nifty shorthand for writing context managers — but how does it work?
Implmeneting contextlib.contextmanager
"What I cannot create, I do not understand." — so let's break it down. Our decorator gets a one-yield generator function, and should return a callable in its stead that, when invoked, returns a context manager guided by its generator object:
>>> def context_manager(gen):
... def wrapper(*args, **kwargs):
... g = gen(*args, **kwargs)
... return GeneratorContext(g)
... return wrapper
This context manager, upon entry, should pass its generator object to next to get it running. When it stops, it means the setup is over, and whatever value it yielded (even if it was just None) is what we should return for the as-clause:
>>> class GeneratorContext:
... def __init__(self, g):
... self.g = g
... def __enter__(self):
... return next(self.g)
Once the context is over, there are two options: on success, we simply resume our object's execution, catching and suppressing its StopIteration (we can even use suppress!). On error, it's a bit trickier — we need to inject the exception into the generator's flow, so it resumes with an error raised. If it handles the error, great; if not, the exception will propagate right back up, as happens in standard context managers. Conveniently, we saw that generator objects have a throw method for precisely this purpose — and it even accepts an error triplet, which is exactly what we have:
>>> class GeneratorContext:
... # same as before
... def __exit__(self, exception, error, traceback):
... with contextlib.suppress(StopIteration):
... if error:
... self.go.throw(exception, error, traceback)
... else:
... next(self.go)
Now see if this makes more sense:
>>> @context_manager
... def example():
... print('before')
... try:
... yield
... finally:
... print('after')
>>> with example():
... print('inside')
before
inside
after
All in all, it's not even that much code; which is nice evidence for our growing mastery in Python, since these measly 16 lines pack a surprising amount of advanced concepts and achieve a pretty sophisticated degree of flow control.