Tag Archives: with

Context managers

I was re-writing the exellent watchedinstall tool and needed to simplify a particularly gnarly chunk of code that required three sub-proceses to be started and then killed after invoking another process. It occurred to me I could make these into context managers.

Previously the code was something like…

start(program1)
try:
    start(program2)
except:
    stop(program1)
    raise

try:
    start(program3)
except:
    stop(program2)
    stop(program1)
    raise

try:
    mainprogram()
finally:
    stop(program3)
    stop(program2)
    stop(program1)

Of course that could have been written with nested try / except / else / finally blocks as well, which I did start with but found not much shorter while almost incomprehensible.

With context managers the whole thing was written as…

# from __future__ import with_statement, Python 2.5

with start(program1):
    with start(program2):
        with start(program3):
            mainprogram()

So much more comprehensible! Here’s the implementation of the context manager (using the contextlib.contextmanager decorator for a triple word score):

import contextlib
import os
import signal
import subprocess


@contextlib.contextmanager
def start(program_args):
    prog = subprocess.Popen(program_args)
    if prog.poll(): # Anything other than None or 0 is BAD
        raise subprocess.CalledProcessError(prog.returncode, program_args[0])

    try:
        yield
    finally:
        if prog.poll() is None:
            os.kill(prog.pid, signal.SIGTERM)

For bonus points I might have used contexlib.nested() to put the three start() calls on one line but then what would I do for the rest of the day?

Django test database runner as a context manager

In my last post I mentioned it might be an idea to wrap up the Django test database setup / teardown in a context manager for use with Python’s with statement. Here’s my first stab, which seems to work.

from contextlib import contextmanager


@contextmanager
def test_db_connection():
    """A context manager for Django's test runner.

    For Python 2.5 you will need
        from __future__ import with_statement
    """

    from django.conf import settings
    from django.test.utils import setup_test_environment, teardown_test_environment
    from django.db import connection

    setup_test_environment()

    settings.DEBUG = False    
    verbosity = 0
    interactive = False

    old_name = settings.DATABASE_NAME
    connection.creation.create_test_db(verbosity, autoclobber=not interactive)

    yield connection

    connection.creation.destroy_test_db(old_name, verbosity)
    teardown_test_environment()

All of this requires Python 2.5 or later.

So with that snippet you could write a test something like so:

import unittest


class MyTestCase(unittest.TestCase):
    def test_myModelTest(self):
        with test_db_connection():
            from myproject.myapp.models import MyModel

            obj = MyModel()
            obj.save()
            self.assert_(obj.pk)

… and just as with Django’s manage.py test command the objects would be created within the test database then destroyed when the with test_db_connection() block is finished.

Everything’s going to be hunky dory.