Factor out admin using context managers

A great thing about Python is how easy it is to throw together a quick automation script. It is empowering to be able to build a tool tailored to whatever task you happen to have.

Many developers write such scripts for their personal use, but sometimes it is desirable to share these: distribute them to your team, to everyone in your organisation or even open source them.

But there are some difficulties associated with this:

• While you were happy to deal with exceptions and tracebacks directly, other users need friendly error handling.
• You don’t know what operating system and Python version your users have.
• Whereas you just wrote logs to the stdout, the user may have complex logging requirements, especially if the script is being run as part of a larger workflow.
• As the user base grows, new corner cases will be uncovered and flawed assumptions exposed.

Photo by Sebastian Herrmann on Unsplash

There are good ways to address all these issues. But there is a real risk that in doing so, you end up writing a lot of additional code and your previously simple script becomes difficult to maintain.

This article will explain how context managers can be used to add administrative logic to your code without obscuring your core business logic.

Example: Talking to a Database

Imagine we have a database that provides an API for retrieving users by ID:

Suppose we want to put some do a bit of logging and error handling when we invoke this function:

The trouble is that now we have a fairly daunting piece of code which in essence does no more than call a function.

Enter Context Managers

A great way to overcome this is to use a context manager. Our code then looks like this:

The context manager does all the admin but then yields control when it’s time to do the core business. The code inside with with statement will be executed at the point that the context manager reaches the yield statement.

Can’t I just write a wrapper function?

Another approach would be to use a wrapper function:

There are two reasons to prefer the context manager.

1. We can parametrise the context manager and re-use it. In the example above, if the database also has a get_product_by_id method, we could use the same context manager and just pass in "product" as the first argument.
2. If we use a wrapper function, the reader has to dive into the internals of that wrapper function to see what it does, in a way defeating the point of abstracting it out. From a readability point of view, a wrapper function says, “get_user_by_id_wrapper is going to give a value to user”, while a context manager says “get_user_by_id is going to give a value to user, under the guidance of this context manager whose workings you may or may not be interested in”.

Conclusion

Computers don’t care whether you use context managers or just put all your administrative logic inline — the behaviour of your program will not change. But code is for humans too.

Any fool can write code that a computer can understand. Good programmers write code that humans can understand. — Martin Fowler

By factoring out your admin cleanly, your code can be production ready yet maintain the simplicity of a quick script, thereby helping the your teammates, the open source community or even your future self understand what’s going on.