The Catch

Python: Tips, Tricks and Idioms - Part 2 - Decorators and Context Managers

by Michael

Two weeks ago I wrote a post called Python: tips, tricks and idioms where I went though a whole lot of features of python. Now however I want to narrow down on just a few, and look at them in more depth. The first is decorators, which I did not cover at all, and the second is context managers which I only gave one example of. Again all the code samples are on gist.github.com.

There is a reason that I put them together; they both have the same goal. They can both help separate what your trying to do (the "business logic") from some extra code added for clean up or performance etc. (the "administrative logic"). So basically it helps package away in a reusable way code that really we don't care too much about.

Decorators

Decorators are really easy to use, and even if you have never seen them before. You could probably guess what is going on, even if not how or why. Take this for example:

Over looking for now exactly what library urlopen comes from... We can guess that the results of the web_lookup() are going to be cached, so that they are not fetched ever time we ask for the same url. So simple, just know we can put @some_decorator before our function and we can use any decorator. But how do we write one?

First we need to understand what the decorator is doing. @cache is actually just syntactic sugar for the following.

So this is important. What cache is, is a function that takes another function as an argument, and returns a new function, that can be used just like the function could be before, but presumably adding in some extra logic. So for our first decorator let us start with something simple, a decorator that squares the result of the function it wraps.

So in this little example every time we call plus() the number will have 1 added it it, but because of our decorator, the result will also be squared.

But there is a problem with this, plus() is no longer the plus() function that we defined, but another function wrapping it. So things like the doc string have gone missing. Things like help(plus) will no longer work. But in the functools library there is a decorator to fix that, functools.wraps(). Always use functools.wraps() when writing a decorator.

But did you notice? wraps() is a decorator that takes arguments. How do we write something like that? It gets a little more involved, but let us start with the code. It will be a function the now raises the number to some power.

So yes three functions, one inside the other. The result now, is that power() is not longer really the decorator, but a function that returns the decorator. See the issue is one of scoping, which we why we have to put the functions inside each other. When _pow() is called, the value of pow, comes from the scope of the power() function that contains it.

So now we know how to write highly reusable function decorators, or do we? There is a problem still, and that is our internal function _square() or _pow() only takes one argument, so that any function it wraps, can only take one argument as well. What we want is to be able to have any number of arguments. So that is where the star operator comes in.

Star operator

The * (star) operator can be used in a function definition so that it can take an arbitrary number of arguments, all of which are collapsed into a single tuple. An example might help.

The * operator can also be used for the reverse case, when we have a iterator, but want to pass that as the arguments to a function. This gets called argument unpacking.

The same basic idea, can also be used for keyword arguments, but for this we use the ** (double star) operator. But instead of getting a list of the arguments, we get a dictionary. We can also use both together. So some examples.

Better Decorators

So now we can go back and change our decorator to be truly generic. Lets do it with the simplest one, we wrote, @square.

Now no matter what arguments the function takes, we will happily just pass them though to the function that we are wrapping.

So let us go back to our web_lookup function, and write first it first with caching, and then write the decorator to see the difference.

That is how it might look, and our problem here is that the caching code is mixed up in with what web_lookup() is supposed to do. Makes it harder to maintain it, reuse it, and also harder if we want to update the way we cache it if we have done something like this all over our code. So our very generic decorator might look like this.

So that can wrap any function with any number of arguments just by putting @cache before it. But I did not write that function myself, I just lifted it right off the Python Decorator Library that has many many examples of decorators you can use.

Context Managers

In the previous post I did a single example of using a context manager, opening a file. It looked like this:

Admittedly the context manager is only a little shorter, and the file will be garbage collected (at least in CPython) but there are other cases where it might be a lot bigger problem if you get it wrong. For example the threading library can also use a context manager.

Which is nice and simple, and you can see from the indent what the critical section is, and also be sure the lock is released.

If you dealing with file like objects that don't have a content manager, the contextlib has the closing context manager. So let us go back an improve our web_lookup() function.

We can also write our own context managers. All that is needed to to used the @contextmanager decorator, and have a function with a yield in it. The yield marks the point at which the context manager stops while the code with in the with statement runs. So the following can be used to time how long it takes to do something.

The try finally in this case is kind of optional, but without it the time would not be printed if there was an exception raised inside the with statement.

We can also do more complicated context managers. The following is something like what was added in Python 3.4. It will take what every is printed to sysout, and put it in a file (or file like object). So for example if we had all our timeit() context managers in your code, and wanted to start putting the results into a log file. Here the yield is followed by a value, which is why we can then use the with ... as syntax.

The last with statement also shows off the use of compound with statements. It is really just the same as putting one with inside another.

Finally it is worth mentioning at least in passing, that any class can be turned into a context manager by adding the __enter__() and __exit__() methods. The code in either will do more or less what the code either side of the yield statement would do.

And that is all for this round, hope you learned something new and interesting. Don't forget to follow me on twitter if you want more python tips, such as when next time I write about sets and dictionaries. In the mean time if you looking for more there is a great book Python Cookbook, Third edition by O'Reilly Media. I have been reading parts of it, and might include a few things I learned from it in my next post. Or if you want something simpler try Learning Python, 5th Edition.

If you found this interesting you might want to