On Clever Usage of Python ‘with’ Clauses

2012-01-21 20:21

As the Python docs state, with statement is used to encapsulate common patterns of tryexceptfinally constructs. It is most often associated with managing external resources in order to ensure that they are properly freed, released, closed, disconnected from, or otherwise cleaned up. While at times it is sufficient to just write the finally block directly, repeated occurrences ask for using this language goodness more consciously, including writing our own context managers for specialized needs.

Those managers – basically a with-enabled, helper objects – are strikingly similar to small local objects involved in the RAII technique from C++. The acronym expands to Resource Acquisition Is Initialization, further emphasizing the resource management part of this pattern. But let us not be constrained by that notion: the usage space of with is much wider.

Interlude on basics

To create our own context manager, we can implement a class with __enter__ and __exit__ methods. The former is called when with block is about to be entered, which is typically associated with obtaining an external resource: file handle, database connection, native memory pool, and so on.

Similarly, once we are about to leave the with block, the __exit__ method will be invoked. If it was an exception that has kicked us out, relevant information will be passed to __exit__ as its arguments; we’ll be able to inspect them if necessary.

More details on both of these methods can be of course found in the documentation, which I encourage to peek at.

Universal cleanup

It is perfectly valid for either of those methods to be a no-op, although there is seldom a point in having only __enter__, Context manager with just __exit__, on the other hand, can be quite useful for executing some finishing action that absolutely must be done. If we also have to account for possible errors, custom with is good alternative to doubly nested try statements that would have to be used instead.

Let’s look at a practical example. On App Engine, there is a Blobstore service which allows to upload & store big files that are infeasible to keep as text or binary data in database. When it comes to uploading them, most of the relevant logic is already present, so we only have to implement the BlobstoreUploadHandler:

  1. from google.appengine.ext.webapp import blobstore_handlers
  3. class UploadHandler(blobstore_handlers.BlobstoreUploadHandler):
  4.     def post(self):
  5.         upload = self.get_uploads()[0]
  6.         # ...do something with uploaded file...
  7.         self.redirect('/upload_complete?status=ok')

The required part here is performing the HTTP redirect at the end, regardless of whether we successfully did something with the uploaded blob (file) or not. But if “doing something” fails, we would like to delete the blob; otherwise, it will just needlessly take space.

Such behavior can be rather easily contained within a specialized context manager with just the __exit__ function:

  1. class blobstore_guard(object):
  2.     def __init__(self, handler):
  3.         self.handler = handler
  4.     def __exit__(self, exc_type, exc_value, traceback):
  5.         if exc_type:
  6.             for _, u in self.handler.get_uploads():
  7.                 u.delete()
  8.             status = 'error'
  9.         else:
  10.             status = 'ok'
  11.         self.handler.redirect('/upload_complete?status=' + status)

If we now put the upload handler’s code inside a with statement with above guard, we will always have a redirect issued, and our blobs will be cleaned up properly in case of error:

  1. class UploadHandler(blobstore_handlers.BlobstoreUploadHandler):
  2.     def post(self):
  3.         with blobstore_guard(self):
  4.             upload = self.get_uploads()[0]
  5.             # ...

The great escape (from with block)

Typical uses of with take advantage of its as part, which is almost always used to capture some kind of resource handle. A file object is a canonical example of such practice:

  1. with open('foo.txt', 'r') as f:
  2.     content = f.read()

If we want to take advantage of this behavior for our own context managers, we shall return a result from the __enter__ method. It will then be given a name (specified after as), and we’ll be able to use it within the with block, as shown above.

But referencing a resource is by no means the only legitimate use for name declared via as. It can actually point to any object, no matter what it is or does.
To illustrate this point, I have came up with rather devious semantical construct which I dubbed a “section”. In some way it’s the closest thing to goto that we can have in Python without resorting to dirty bytecode hacks. It doesn’t support actual jumps (i.e. moving instruction pointer into arbitrary place), but allows to prematurely “bail” out of the with scope, if desired. Hence, it is analogous to the break label mechanism from Java, while also serving similar purpose to infamous do-while(false) construct:

  1. with section() as s:
  2.     for i in xrange(5):
  3.         for j in xrange(5):
  4.             if i * j > 20: s.bail()
  6. print "This will be executed after bail()"

Of course, there aren’t many uses for such a thing, and I would prefer not to maintain code that actually employs this trick. But it’s entirely sufficient as an example – albeit somewhat silly. Here’s how it could be implemented:

  1. class section(object):
  2.     class Bail(Exception):
  3.         def __init__(self, s):
  4.             self.section = s
  5.     def __enter__(self):
  6.         return self
  7.     def __exit__(self, exc_type, exc_value, traceback):
  8.         return exc_type == Bail and exc_value.section == self
  9.     def bail(self):
  10.         raise Bail(self)

As you can see, we’re using a special exception as means for “bailing out” of the section. The interesting bits happen mostly inside __exit__, where we examine the type of exception (if any) and the section it came from (to account for possible nesting). We are then returning a boolean value, indicating whether the exception should be squelched.
The __enter__ function is also notable, since it’s quite common to just return self from it. This way, the context manager itself becomes the as-object. Not accidentally, this is also the case with the file type, which is what allows for open function to be used both normally and in withas clause. For section, such behavior is not required, but __enter__ returning self is simply the most straightforward solution.

More ideas

Looking at slightly bigger picture, we can see that with statements create something which might be called a programmable scope. Thanks to __enter__ and __exit__ methods, we are able to capture the events of going in and out of these blocks, and act upon them.

I suspect that when used cleverly, this may result in some pretty powerful and impressive programming techniques, not really related to any sort of resource management. I tried to show a glimpse, but I believe there are many more fruits to be picked from that tree – including:

  • Intrinsic profiling. For one, it wouldn’t be hard to implement something like a frame rate counter for logic contained in a with block. This could be further expanded into measurements of varying granularity, thanks to with blocks’ nesting.
  • Parsing. Done by hand, parsing involves rather cumbersome state management. We have to save current position within parsed text, and restore it should we need to try a different variant. This looks very much like a pattern to be encapsulated in specialized with blocks. If done correctly, it could result in nice and neat API that wouldn’t hide the control flow of parsing process – unlike most (all?) libraries that rely on explicit specification of grammar.
  • Testing utilities. As one almost obvious example, with could be used to check whether test code throws an expected exception. It would be much more convenient (and readable) than assertRaises method from standard unittest package.
  • “Important” assignments. As purely syntactic device, we could use with to put emphasis on certain notable assignments in our code. This is vaguely similar to letin or where clauses from functional languages, but of course it doesn’t have any purity implications. The downside: we would often need to use a “fake” context manager if object being assigned isn’t already one:
    1. class var(object):
    2.     def __init__(self, obj): self.obj = obj
    3.     def __enter__(self): return self.obj
    4.     def __exit__(self, et, ev, t): pass

This list is of course nowhere near complete. I encourage to look at the with statement in similarly unorthodox way, for there are likely many semantic gems (or at least polished rocks ;]) to be uncovered here.

Be Sociable, Share!
Be Sociable, Share!

One comment for post “On Clever Usage of Python ‘with’ Clauses”.
Comments are disabled.

© 2023 Karol Kuczmarski "Xion". Layout by Urszulka. Powered by WordPress with QuickLaTeX.com.