Python with Statements Using Contextlib

by James Johnson

The with statement in Python is one of my favorite Python features. This post discusses an easy way to manage contexts using contextlib as well as a deeper dive if you were to do it manually.

Python Contexts using with

The Python with statement was first introduced in pep-343 in order to:

make it possible to factor out standard uses of try/finally statements.

File Open Example

One of the most common uses of the with statement is when opening files:

with open("/tmp/test.txt", "wb") as f:
    f.write("new contents")

When the with block exits, the file will be closed, even if exceptions occur within the block.

This essentially replaces a try/finally pattern (see also my post on python exception handling):

try:
    f = open("/tmp/test.txt", "wb")
    f.write("new contents")
finally:`
    f.close()

Using contextlib

The contextlib package is included as part of Python’s standard library. contextlib provides three functions:

  • contextmanager - a decorator that defines a factory function that can be used with a with statement
  • nested - This has been deprecated in favor of compound with statements
  • closing - a context manager that closes that supplied argument after the with block is executed

The contextmanager decorator is usually what you will be using when creating your own context generators.

For example, to define a function that will clone a github repository into a temporary directory and cleanup once the with block has executed, one could do something like:

from contextlib import contextmanager
import os
from sh import git
import shutil
import tempfile

@contextmanager
def github_clone(project):
    project_url = "https://github.com/" + project
    project_name = project_url.rsplit("/",1)[-1]
    tmpdir = tempfile.mkdtemp()

    project_path = os.path.join(tmpdir, project_name)
    git.clone(project_url, project_path)
    git_ = git.bake(
        "--work-tree", tmpdir,
        "--git-dir", os.path.join(project_path, ".git")
    )

    try:
        yield git_,project_path
    finally:
        shutil.rmtree(tmpdir)

with github_clone("d0c-s4vage/pfp") as info:
    git_,tmpdir = info
    tags = git_.tag().split("\n")
    print("pfp tags: " + ",".join(tags))

Notice that we must explicitly wrap the yield statement in try/finally blocks to ensure that our cleanup code is run despite any exceptions that may occur.

Manually Creating Context Managers

Context managers require two functions to be defined in order to be used in a with statement:

  • __enter__ - Code that should be executed prior to executing the with block
  • __exit__ - Code that should always be executed after executing the with block

Manual Context Manager Example

import tempfile
import os
from sh import git
import shutil

class github_clone(object):
    def __init__(self, project):
        self.tmpdir = tempfile.mkdtemp()

        project_url = "https://github.com/" + project
        project_name = project_url.rsplit("/",1)[-1]

        tmpdir = tempfile.mkdtemp()
        self.project_path = os.path.join(tmpdir, project_name)
        git.clone(project_url, self.project_path)
        self.git_ = git.bake(
            "--work-tree", tmpdir,
            "--git-dir", os.path.join(self.project_path, ".git")
        )

    def __enter__(self):
        return (self.git_,self.project_path)

    def __exit__(self, exc_type, exc_val, exc_tb):
        shutil.rmtree(self.tmpdir)
        # explicitly raise any exceptions after cleaning up
        return False

with github_clone("d0c-s4vage/pfp") as info:
    git_,tmpdir = info
    tags = git_.tag().split("\n")
    print("pfp tags: " + ",".join(tags))

Notice that the __enter__ function takes no parameters, and the __exit__ function takes the exception info for any exceptions that occurred while executing the with block. If no exceptions occurred while executing the with block, the exception infos will be None.

Also notice the return value of the __exit__ function. If True is returned, any exceptions that occurred will be suppressed. If False is returned, the exception will be reraised after the __exit__ function returns.