Friday, November 16, 2012

Context Managers Are Your Friend

Anytime there is built in support for functionality that reduces the potential for bugs (and number of lines of code), while simultaneously improving code legibility, that functionality should be used with reckless abandon. Similar to function decorators that can be used to elegantly wrap functions with additional functionality, context managers allow often repeated blocks of code to be elegantly written.

Let's take a common database transaction use-case as an example.
try:
    db.insert(some_pbject)
except:
    db.rollback()
else:
    db.commit()

... do some stuff ...

try:
    db.insert(some_other_pbject)
except:
    db.rollback()
else:
    db.commit()

As you can see, there is nothing special about the (pseudo) code above, as it simply is some code to insert some objects into a database. There isn't really anything wrong with above code, but using context managers we can do this in a much more elegant way:
from contextlib import contextmanager

@contextmanager
def transaction(db):
    try:
        yield
    except:
        db.rollback()
    else:
        db.commit()

with transaction(db):
    db.insert(some_pbject)

... do some stuff ...

with transaction(db):
    db.insert(some_other_pbject)

Ok, ok, that's about the same number of lines of code than the original, but now, every time you want to commit something to the database, you never need to write that try-except-else block to make sure that things are either committed or rolled back - you merely need to use the with transaction(db) statement and you're good to go! And since you only wrote the commit logic once, if there is a bug, it only needs to be fixed in one location. Similarly, if after you've written a bunch more code you decide that you want to log every time there is an exception, you simply have to add the logging code in the except block in your transaction function.

Now that you've gotten a taste of the power of context managers, you will probably want to know more about how they actually work. As there is a great explanation here and here, I'm not going to cover that in this post. You may also check out the Python docs on the @contextmanager.

No comments:

Post a Comment