Understanding Python context managers by reading Django source code

Django comes with a bunch of useful context managers. We will read their source code to find what context managers can do and how to implement them including some best parctices.

The three I use most are

  • transactions.atomic - To get a atomic transaction block
  • TestCase.settings - To change settings during a test run
  • connection.cursor - TO get a raw cursor

connection.cursor Is generally implemented in the actual DB backends such a psycopg2, so we will focus on transactions.atomic, TestCase.settings and a few other contextmanagers.

What is a context manager?

Context managers are a code patterns for

  • Step 1: Do something
  • Step 2: Do something else
  • Step 3: Final step, this step must be guaranteed to run.

For example when you say

with transaction.atomic():
    # This code executes inside a transaction.
    do_more_stuff()

What you really want is:

  • create a savepoint
  • do_more_stuff()
  • Commit or rollback the savepoint

Similarly, when you say (Inside a django.test.TestCase)

with self.settings(LOGIN_URL='/other/login/'):
    response = self.client.get('/sekrit/')

What you want is

  • Change settings to LOGIN_URL=’/other/login/’
  • response = self.client.get('/sekrit/'), assert something with on response with the changed setting.
  • Change settings back to what existed at start.

A context manager povides a clean api to enforce this three step workflow.

Some non-Django context managers

The most common context manager is

with open('alice-in-wonderland.txt', 'rw') as infile:
    line = infile.readlines()
    do_something_more()

If you did not have open contextmanager, you would need to do the below everytime, because you need to ensure do_something_more() is called.

try:
    infile = open('alice-in-wonderland.txt', 'r')
    line = infile.readlines()
    do_something_more()
finally:
    infile.close()

Another common use is

a_lock = threading.Lock()

with a_lock:
    do_something_more()

And without a context manager, this would have been.

a_lock.acquire()
try:
    do_something_more()
finally:
    a_lock.release()

So at a high level, context managers are syntactic sugar for ``try: … finally …`` block. This is important, so I will repeat context managers are syntactic sugar for ``try: … finally …`` block

Implementing context managers

Context managers can be implemented as a class with two required methods and one optional __init__

  • __enter__: what to do when the context starts
  • __exit__: what to do when the context ends
  • __init__: if your context manager requires arguments

Alternatively, you can use contextlib.contextmanager with yield statements to get a context manager. We will see an example in the next section.

A simple Django context manager

In django/tests/backends/mysql/tests.py, Django implements a very simple context manager.

@contextmanager
def get_connection():
    new_connection = connection.copy()
    yield new_connection
    new_connection.close()

And then uses it like this:

def test_setting_isolation_level(self):
    with get_connection() as new_connection:
        new_connection.settings_dict['OPTIONS']['isolation_level'] = self.other_isolation_level
        self.assertEqual(
            self.get_isolation_level(new_connection),
            self.isolation_values[self.other_isolation_level]
        )

There is some code here which doesn’t immediately concern us, let us just focus on with get_connection() as new_connection:

Using @contextmanager, here is what happened:

  • The part before yield new_connection = connection.copy() handles the context setup.
  • The yield new_connection part allows using new_connection as as new_connection.
  • The part after yield new_connection.close() handle context teardown.

Lets look at the TestCase.settings next, which uses the __enter__ - __exit__ protocol.

Implementing Testcase.settings

Testcase.settings is implemented as

def settings(self, **kwargs):
    """
    A context manager that temporarily sets a setting and reverts to the
    original value when exiting the context.
    """
    return override_settings(**kwargs)

There is a bit of class hierarchy to jup through which takes us from

Testcase.settings → override_settings → TestContextDecorator

Skipping the part we don’t care about, we get

class TestContextDecorator:
    # ...
    def enable(self):
        raise NotImplementedError

    def disable(self):
        raise NotImplementedError

    def __enter__(self):
        return self.enable()

    def __exit__(self, exc_type, exc_value, traceback):
        self.disable()

And then override_settings implements .enable and .disable

class override_settings(TestContextDecorator):
    # ...
    def enable(self):
        # Keep this code at the beginning to leave the settings unchanged
        # in case it raises an exception because INSTALLED_APPS is invalid.
        if 'INSTALLED_APPS' in self.options:
            try:
                apps.set_installed_apps(self.options['INSTALLED_APPS'])
            except Exception:
                apps.unset_installed_apps()
                raise
        override = UserSettingsHolder(settings._wrapped)
        for key, new_value in self.options.items():
            setattr(override, key, new_value)
        self.wrapped = settings._wrapped
        settings._wrapped = override
        for key, new_value in self.options.items():
            setting_changed.send(sender=settings._wrapped.__class__,
                                 setting=key, value=new_value, enter=True)

    def disable(self):
        if 'INSTALLED_APPS' in self.options:
            apps.unset_installed_apps()
        settings._wrapped = self.wrapped
        del self.wrapped
        for key in self.options:
            new_value = getattr(settings, key, None)
            setting_changed.send(sender=settings._wrapped.__class__,
                                 setting=key, value=new_value, enter=False)

There is a lot of boiler plate here which is interesting, but skipping the state management we see

class override_settings(TestContextDecorator):
    # ...
    def enable(self):
        # ...
        # This gets called by __enter__
        for key, new_value in self.options.items():
            setattr(override, key, new_value)
        self.wrapped = settings._wrapped
        settings._wrapped = override
        for key, new_value in self.options.items():
            setting_changed.send(sender=settings._wrapped.__class__,
                                 setting=key, value=new_value, enter=True)

    def disable(self):
        # ...
        # This gets called by __exit__
        for key in self.options:
            new_value = getattr(settings, key, None)
            setting_changed.send(sender=settings._wrapped.__class__,
                                 setting=key, value=new_value, enter=False)

Implmenting context manager to also be used as a decorator.

When you can say with transaction.atomic():, you can get the same effect by using it as a decorator.

@transaction.atomic
def do_something():
    # this must run in a transaction
    # ...

Implmenting a context manager to also be used as a decorator is a common pattern and Django does the same with atomic. contextlib.ContextDecorator makes this straightforward.

# class Atomic is implemented later
def atomic(using=None, savepoint=True):
    # Bare decorator: @atomic -- although the first argument is called
    # `using`, it's actually the function being decorated.
    if callable(using):
        return Atomic(DEFAULT_DB_ALIAS, savepoint)(using)
    # Decorator: @atomic(...) or context manager: with atomic(...): ...
    else:
        return Atomic(using, savepoint)

class Atomic(ContextDecorator):
    # There is a lot of complicated corner cases and error handling.
    # See the gory details in django/django/db/transaction.py
    def __init__(self, using, savepoint):
        self.using = using
        self.savepoint = savepoint

    def __enter__(self):
        connection = get_connection(self.using)
        # ...
        # sid = connection.savepoint()
        # connection.savepoint_ids.append(sid)

    def __exit__(self, exc_type, exc_value, traceback):
        # Skip the gory details
        # ...
        sid = connection.savepoint_ids.pop()
        if sid is not None:
            try:
                connection.savepoint_commit(sid)
            except DatabaseError:
                connection.savepoint_rollback(sid)

Final thoughts

Context managers provide a simple API for a powerfulo construct. Even though they are merely syntactic sugar, they make for an itutive API and in conjunction with the contextlib module are easy to implement.