The curious case of Python's context manager
Python's context managers are great for resource management and stopping the propagation of leaked abstractions. You've probably used it while opening a file or a database connection. Usually it starts with a with statement like this:
In the above case, file.txt gets automatically closed when the execution flow goes out of the scope. This is equivalent to writing:
Writing custom context managers
To write a custom context manager, you need to create a class that includes the enter and exit methods. Let's recreate a custom context manager that will execute the same workflow as above.
You can use the above class just like a regular context manager.
From generators to context managers
Creating context managers by writing a class with enter and exit methods isn't difficult. However, you can achieve better brevity by defining them using contextlib.contextmanager decorator. This decorator converts a generator function into a context manager. The blueprint for creating context manager decorators goes something like this:
When you use the context manager with the with statement:
It roughly translates to:
The setup code goes before the try..finally block. Notice the point where the generator yields. This is where the code block nested in the with statement gets executed. After the completion of the code block, the generator is then resumed. If an unhandled exception occurs in the block, it's re-raised inside the generator at the point where the yield occurred and then the finally block is executed. If no unhandled exception occurs, the code gracefully proceeds to the finally block where you run your cleanup code.
Let's implement the same CustomFileOpen context manager with contextmanager decorator.
Now use it just like before:
Writing context managers as decorators
You can use context managers as decorators also. To do so, while defining the class, you have to inherit from contextlib.ContextDecorator class. Let's make a RunTime decorator that'll be applied on a file-opening function. The decorator will:
- Print a user provided description of the function.
- Print the time it takes to run the function.
You can use the decorator like this:
Using the function like this should return:
You can also create the same decorator via contextlib.contextmanager decorator.
Nesting contexts
You can nest multiple context managers to manage resources simultaneously. Consider the following dummy manager:
Notice the order they're closed. Context managers are treated as a stack, and should be exited in reverse order in which they're entered. If an exception occurs, this order matters, as any context manager could suppress the exception, at which point the remaining managers will not even get notified of this. The exit method is also permitted to raise a different exception, and other context managers then should be able to handle that new exception.
Combining multiple context managers
You can combine multiple context managers too. Let's consider these two managers.
Now combine these two using the decorator syntax. The following function takes the above define managers a and b and returns a combined context manager ab.
This can be used as:
If you have variable numbers of context managers and you want to combine them gracefully, contextlib.ExitStack is here to help. Let's rewrite context manager ab using ExitStack. This function takes the individual context managers and their arguments as tuples and returns the combined manager.
ExitStack can be also used in cases where you want to manage multiple resources gracefully. For example, suppose, you need to create a list from the contents of multiple files in a directory. Let's see, how you can do so while avoiding accidental memory leakage with robust resource management.
Using context managers to create SQLAlchemy session
If you are familiar with SQLALchemy, Python's SQL toolkit and Object Relational Mapper, then you probably know the usage of Session to run a query. A Session basically turns any query into a transaction and make it atomic. Context managers can help you write a transaction session in a very elegant way. A basic querying workflow in SQLAlchemy may look like this:
The excerpt above creates an in memory SQLite connection and a session_scope function with context manager. The session_scope function takes care of committing and rolling back in case of exception automatically. The session_scope function can be used to run queries in the following way:
Abstract away exception handling monstrosity with context managers
This is my absolute favorite use case of context managers. Suppose you want to write a function but want the exception handling logic out of the way. Exception handling logics with sophisticated logging can often obfuscate the core logic of your function. You can write a decorator type context manager that will handle the exceptions for you and decouple these additional code from your main logic. Let's write a decorator that will handle ZeroDivisionError and TypeError simultaneously.
Now use this in a function where these exceptions occur.
You can see that the errhandler decorator is doing the heavylifting for you. Pretty neat, huh?
The following one is a more sophisticated example of using context manager to decouple your error handling monstrosity from the main logic. It also hides the elaborate logging logic from the main method.
This will return
Persistent parameters across HTTP requests with context managers
Another great use case for context managers is making parameters persistent across multiple HTTP requests. Python's requests library has a Session object that will let you easily achieve this. So, if you're making several requests to the same host, the underlying TCP connection will be reused, which can result in a significant performance increase. The following example is taken directly from the official docs of the requests library. Let's persist some cookies across requests.
This should show:
Remarks
To avoid redundencies, I have purposefully excluded examples of nested with statements and now deprecated contextlib.nested function to create nested context managers.
Further reading
Discussion in the ATmosphere