Deferring in Python

I move a lot between Python, Java, Rust and Go lately. All of them have strong and weak points, and require approaching problems differently. I often find myself missing features from other languages when developing with any of them. Sometimes it is because I am using the wrong approach, sometimes it is just because the language is missing features. One of this features is Go’s defer when working with Python.

Quoting The Go Blog:

A defer statement pushes a function call onto a list. The list of saved calls is executed after the surrounding function returns. Defer is commonly used to simplify functions that perform various clean-up actions.

I would summarize defer in three points:

1. It is a stack (LIFO) of function callbacks with arguments.
2. Every function is called after the function that created it ends, regardless of the reason of the exit (success or error).
3. Every function is called regardless of the success of the rest of the functions.

I’ve encountered many cases where tests have complex cleanups, with multiple statements that must execute regardless of the success of the previous statements; which often ends up looking like either several subsequent try/except blocks, or even worse, several nested try/except blocks. An example of this is:

def test_insert():
    schema = create_schema()
    table = create_table()
    try:
        # Do the actual test
    finally:
        try:
            table.drop()
        except Exception as e:
            logger.error(e)
        try:
            schema.drop()
        except Exception as e:
            logger.error(e)

Go’s defer would have several advantages here:

1. No nesting, which would make the code more readable.
2. The defer is usually done immediately after the object to be deleted is created.

By introducing a cleanup fixture that mimics the behavior of Go’s defer in Python, we can simplify the code above to look like this:

from cleanup import cleanup

def test_insert(cleanup):
    schema = create_schema()
    cleanup(schema.drop)
    table = create_table()
    cleanup(table.drop)
    # Do the actual test

This code is more concise and readable. The fixture is implemented throught Python’s native ExitStack, which behave the same way as Go’s defer but within their with block. We can manage it’s lifetime by yielding it:

import contextlib
import logging
import pytest

logger = logging.getLogger(__name__)

@pytest.fixture()
def cleanup(request):
    try:
        with contextlib.ExitStack() as exit_stack:
            yield exit_stack.callback
    except Exception as exception: # Log the last error raised in the cleanup stack
        logger.error(f"Error during cleanup of {request.node.name}: {exception}")

Wrapping up: Exit Stacks are awesome, and can be used outside tests too. Happy hacking!