Python Core Utility - `contextlib.contextmanager`

Intro

In this post, we will look at the contextmanager function from the contextlib module.

The contextlib.contextmanager function is used as a decorator that allows you to create context managers using a generator function instead of writing a class with __enter__ and __exit__ methods. It's a cleaner, more pythonic way to handle resource management, setup/teardown operations, and temporary state changes.

Context managers in a nutshell

You might have come across code like:

with open(file_name, "w") as f:
    do_something(f)

This code block opens the file_name file in write-mode, stores the opened file handle in the variable f. This f can then be used in the scope of the with block.

The thing that is passed to with (in this case, the result of open(..)) is a context manager.

Context managers have methods that run before and after the with block execution. These methods are, respectively, __enter__() and __exit__(). That means, any traditional class that defines these methods can be used as a context manager (see example below). But, with the @contextmanager decorator, you can achieve the same result with a simple generator function. And that's what we are covering in this post here.

However, for the sake of completeness, here's how you can create your own context managers using a class. This example creates a requests.Session to load the auth token into headers. Then, upon entering the with block it first calls the /login endpoint, and upon exiting the block it calls the /logout endpoint. Inside the block, you can do your thing, call different endpoints, without worrying about login or logout.

import requests

class SessionManager:
    def __init__(self, base_url, token):
        self.base_url = base_url.rstrip("/")
        self.token = token

    def __enter__(self):
        self.session = requests.Session()
        self.session.headers.update({"Authorization": f"Bearer {self.token}"})
        response = self.session.post(f"{self.base_url}/login")
        response.raise_for_status()
        return self.session

    def __exit__(self, exc_type, exc_value, traceback):
        self.session.post(f"{self.base_url}/logout")

with SessionManager("https://my-base-url.com", "my-fake-token") as session:
    refresh_reports(session)
    download_reports(session)
    email_reports(session)

A bit more about the __exit__ method - it accepts three arguments (all optional):

• the exception class
• the exception instance
• the traceback object

If an exception occurs in the with block, these parameters contain the exception details. If you have set __exit__ to return True, the exception raised inside the with block is suppressed, and if it returns False (or None), the exception propagates to outside of the block.

From the standard docs:

If an exception is supplied, and the method wishes to suppress the exception (i.e., prevent it from being propagated), it should return a true value. Otherwise, the exception will be processed normally upon exit from this method.

Note that exit() methods should not reraise the passed-in exception; this is the caller’s responsibility.

Use cases

The primary use cases for using contextmanager could be:

• Resource management - Opening and closing files, database connections, network sockets, wrapping API calls between generic login and logout, etc.
• Setup and teardown - Setting up test fixtures, temporary configurations, mocking
• Temporary state changes - Changing working directory, modifying environment variables, suppressing output
• Locking and synchronization - Acquiring and releasing locks
• Transaction management - Begin/commit/rollback database transactions

Usage

You probably already know the traditional, class-based approach to creating context managers (using the __enter__ and __exit__ methods). If not, I have covered it briefly in the Context managers in a nutshell part in Intro section above.

Now let's see how the @contextmanager decorator can be used:

Consider the same example as given before - a wrapper that calls the /login and /logout endpoints via requests.Session before and after the with block:

from contextlib import contextmanager
import requests

@contextmanager
def session_manager(base_url, token):
    try:
        session = requests.Session()
        session.headers.update({"Authorization": f"Bearer {token}"})
        response = session.post(f"{base_url}/login")
        response.raise_for_status()
        yield session
    finally:
        session.post(f"{base_url}/logout")

with session_manager("https://my-base-url.com", "my-fake-token") as session:
    refresh_reports(session)
    download_reports(session)
    email_reports(session)

Here's how it works:

1. Code before yield runs when entering the with block (like __enter__)
2. The yield statement returns (more correctly, yields) what gets assigned to the with variable (session in this case)
3. Code after yield runs when exiting the with block (like __exit__)
4. In the above example, we could also get rid of the try..finally completely and have the /logout call immediately after the yield statement. But putting the logout (or any closing functionality as such) in finally ensures that it gets called no matter what.

One might get inclined in the favour of the class-based implementation over this decorator based one. And that's perfectly fine! I myself have used the class-based method several times - it just depends on how complex your use case is, and how much more work you wish to do around __enter__ and __exit__.

Some other common use cases of the decorator method can include:

• Timing code execution:
  

  from contextlib import contextmanager
  import time
  
  @contextmanager
  def timer(label):
      start = time.time()
      try:
          yield
      finally:
          end = time.time()
          print(f"{label}: {end - start:.4f} seconds")
  
  with timer("Processing"):
      do_something()
  
  # [Out]: Processing: 0.5001 seconds
  

• Temporary directory change:
  

  import os
  from contextlib import contextmanager
  
  @contextmanager
  def change_dir(new_dir):
      old_dir = os.getcwd()
      os.chdir(new_dir)
      print(f"Changing dir from {old_dir} to {new_dir}")
      try:
          yield
      finally:
          os.chdir(old_dir)
          print(f"Changing dir back to {old_dir} from {new_dir}")
  
  with change_dir("/tmp"):
      print(f"Inside {os.getcwd()}")
  
  # [Out]:
  # Changing dir from /home/shodhan/projects/blog to /tmp
  # Inside /tmp
  # Changing dir back to /home/shodhan/projects/blog from /tmp
  

• Suppressing stdout:
  

  import sys
  from contextlib import contextmanager
  from io import StringIO
  
  @contextmanager
  def suppress_stdout():
      old_stdout = sys.stdout
      sys.stdout = StringIO()
      try:
          yield
      finally:
          sys.stdout = old_stdout
  
  print("This prints")
  with suppress_stdout():
      print("This doesn't print")
  print("This prints again")
  # [Out]:
  # This prints
  # This prints again
  

Real-life scenarios

Here are some practical examples where @contextmanager shines.

Database transaction management

One of the most common uses - ensuring transactions are committed or rolled back:

from contextlib import contextmanager
import sqlite3

@contextmanager
def db_transaction(db_path, commit=False):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    try:
        yield cursor
        conn.commit()
    except Exception as e:
        conn.rollback()
        raise
    finally:
        conn.close()

with db_transaction("mydb.db") as cursor:
    cursor.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
    cursor.execute("INSERT INTO users (name) VALUES (?)", ("Bob",))

Note the raise statement here - Above we said that we do not raise from the context manager, but this is different - here we are raising the same exception that we caught explicitly, because we wanted to do another thing (conn.rollback) before it bubbled up.

Temporary environment variables

When testing code that depends on environment variables:

import os
from contextlib import contextmanager

@contextmanager
def temp_env_var(key, new_value):
    old_value = os.environ.get(key)
    os.environ[key] = new_value
    try:
        yield
    finally:
        if old_value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = old_value

print(os.environ.get("DEBUG"))
# [Out]: None

with temp_env_var("DEBUG", "true"):
    print(os.environ.get("DEBUG"))
    # [Out]: true

print(os.environ.get("DEBUG"))
# [Out]: None

A special context manager (nullcontext) to parametrize pytest fixtures

This was a sweet surprise when I found it! The nullcontext is another function of contextlib, and is a special case of contextmanager, in the sense that it is a context manager that does nothing! I will write in-depth about it in a separate post. For now, let's see how it can help in the pytest parametrized fixtures.

When writing pytest parametrized fixtures, where some cases should raise exceptions and others shouldn't, you can use nullcontext for cases that should not raise exceptions. It can be used here to avoid writing an if..else block for cases that raise exceptions and cases that don't. This keeps the code quite elegant!

import pytest
from contextlib import nullcontext

def divide(a, b):
    """Division function that raises on zero."""
    return a / b

@pytest.mark.parametrize(
    "a, b, expected, raised_exception",
    [
        (10, 2, 5, nullcontext()),
        (20, 4, 5, nullcontext()),
        (10, 0, None, pytest.raises(ZeroDivisionError)),
        (0, 0, None, pytest.raises(ZeroDivisionError)),
    ],
)
def test_divide(a, b, expected, raised_exception):
    with raised_exception:
        result = divide(a, b)
        if expected is not None:
            assert result == expected

Managing multiple resources

When you need to manage multiple related resources, you can call the same context manager multiple times. Consider the example of temporary environment variables from above, we can call it multiple times like this:

with (
    temp_env_var("DEBUG", "true"),
    temp_env_var("RUN_ENV", "prod"),
):
    do_something()

Async context managers

All the examples mentioned here are synchronous, including the ones opening database connection, or making the session requests. But contextlib also provides another function - asynccontextmanager for async operations. Although this is technically a different function than what this post is about, I'd like to include it here as I feel it is closely related and will also make sense to you.

This works similarly to the regular, "sync" context manager, except that instead of calling it with with, you must call it with async with block. And of course, the decorator function should also be async:

Code copied from the standard docs

from contextlib import asynccontextmanager

@asynccontextmanager
async def get_connection():
    conn = await acquire_db_connection()
    try:
        yield conn
    finally:
        await release_db_connection(conn)

async def get_all_users():
    async with get_connection() as aconn:
        return aconn.query("SELECT ...")

Nuances

Must use try/finally if you need cleanup

If you have cleanup code, it MUST be in a finally block to ensure it runs even if an exception occurs:

from contextlib import contextmanager

# WRONG - cleanup might not happen
@contextmanager
def bad_context():
    resource = acquire_resource()
    yield resource
    release_resource(resource)  # might not run if exception occurs!

# CORRECT - cleanup always happens
@contextmanager
def good_context():
    resource = acquire_resource()
    try:
        yield resource
    finally:
        release_resource(resource)  # always runs

Only one yield allowed

A context manager can only have one yield statement. If you have multiple yields, only the first one is used:

from contextlib import contextmanager

@contextmanager
def bad_multiple_yields():
    yield 1
    yield 2  # this never happens

with bad_multiple_yields() as value:
    print(value)
# [Out]: 1

Exceptions before yield

If an exception occurs before the yield, it propagates immediately and __exit__ is never called:

from contextlib import contextmanager

@contextmanager
def failing_context():
    print("Setting up")
    raise ValueError("Setup failed")
    yield  # never reached
    print("Cleaning up")  # never reached

try:
    with failing_context() as ctx:
        print("In context")  # never reached
except ValueError as e:
    print(f"Caught: {e}")

# [Out]:
# Setting up
# ValueError: Setup failed

This is expected behavior - if setup fails, there's nothing to clean up.

This is the reason why I put response.raise_for_status() in the session manager example above, so that if there is any error during login, it must not proceed ahead!

Cannot suppress exceptions by default

Unlike class-based context managers where you can return True from __exit__ to suppress exceptions, with @contextmanager you need to catch and handle exceptions explicitly:

from contextlib import contextmanager

# This does NOT suppress exceptions
@contextmanager
def does_not_suppress():
    try:
        yield
    finally:
        pass  # exception still propagates

# To suppress, you must catch explicitly
@contextmanager
def suppress_errors():
    try:
        yield
    except Exception as e:
        print(f"Suppressed: {e}")
        # by not re-raising, we suppress the exception

with suppress_errors():
    raise ValueError("This is suppressed")
    print("This won't print")

print("Program continues")
# [Out]:
# Suppressed: This is suppressed
# Program continues

Nested context managers can be tricky

Be careful when nesting context managers - exceptions in inner context managers can affect outer ones:

from contextlib import contextmanager

@contextmanager
def outer():
    print("Outer enter")
    try:
        yield "outer"
    finally:
        print("Outer exit")

@contextmanager
def inner():
    print("Inner enter")
    try:
        yield "inner"
    finally:
        print("Inner exit")

# nested usage
with outer():
    with inner():
        print("Inside inner")
        raise ValueError("Error!")

# [Out]:
# Outer enter
# Inner enter
# Inside inner
# Inner exit
# Outer exit
# ValueError: Error!

Also, looking at the print statements, you can see that the cleanup order should be inside-out (like nested try/finally blocks).

Reusability

Context managers created with @contextmanager are reusable - you can use them multiple times:

from contextlib import contextmanager

@contextmanager
def my_context():
    print("Enter")
    yield
    print("Exit")

ctx = my_context()

with ctx:
    print("First use")
# [Out]:
# Enter
# First use
# Exit

with ctx:
    print("Second use")
# [Out]:
# Enter
# Second use
# Exit

However, be aware that the generator function runs fresh each time, so any state is reset.

Other types of context managers

If you check the standard docs, you will see that there are several different types of context managers, made for special cases that are believed to be quite frequent.

I would really like to write about them, but perhaps in a separate blog post. I will update this space once I have that post published.

🙏 Thanks for taking out the time to read this! And as always, please do not hesitate to give feedback on my posts!