ExitStack in Python

Over the years, I've used Python's contextlib.ExitStack in a few interesting ways. The official ExitStack documentation advertises it as a way to manage multiple context managers and has a couple of examples of how to leverage it. However, neither in the docs nor in GitHub code search could I find examples of some of the maybe unusual ways I've used it in the past. So, I thought I'd document them here.

Enforcing request level transaction

While consuming APIs, it's important to handle errors in a way that prevents database state corruption. In the following example, I'm making two POST requests to an API and rolling back to the original state if any one of them fails:

Running this will print the following output:

Here, the group_create function makes two calls to POST httpbin.org/post endpoint and the maybe_rollback function deletes the created record if any one of the two requests fails. In the main function, I've used the ExitStack.callback method to register the maybe_rollback callback. If you change the expected_status_code in the maybe_rollback function to something like HTTPStatus.FORBIDDEN, you'll be able to see the cleanup callbacks in action:

Invoking conditional event hooks

The same strategy used in the previous section can be applied to invoke event hooks conditionally. For example, let's say you want to run a callback function when some event function executes. However, you want only a particular type of callback function to be executed depending on the state of your conditionals or code path. I've found the following pattern useful in this case:

Here the .on_failure hook will only be called if there's an error in your execution path raises an exception.

Avoiding nested context structure

It can get ugly pretty quickly when you start using multiple nested context managers. For example, if you need to open two files and copy content from one file to the other, you'd typically start two nested context managers and transfer the content like this:

ExitStack can help you get away with only one level of nesting here. Here's a complete example:

This example creates two in-memory temporary file instances with tempfile.SpooledTemporaryFile. The SpooledTemporaryFile can be used as a context manager. However, instead of nesting the two instances, I'm using ExitStack.enter_context to enter into the context manager without explicitly using the with statement. This .enter_context method ensures that the exit method of the respective context managers will be called properly at the end of the main() function run.

Then in the body of the ExitStack, we're writing some content to the first in-memory file and then copying the content to the other in-memory file. If we had to open and manage even more context managers, in this way, we'd be able to that without crating any additional nestings.

Applying multiple patches as context managers

Python's unittest.mock.patch can be used as both decorators and context managers. For granular patching and unpatching during tests, the context manager approach gives you more control than its decorator counterpart. In this case, ExitStack can help you avoid multiple nestings just like in the previous section:

Running the above snippet with pytest will reveal that the test passes without any error:

Here, I'm making GET and POST requests with the httpx library and in the test_main function, the httpx.get and httpx.post callable are patched with the patch context manager. However, ExitStack allows me here to do it without creating additional nested with blocks.