Python Context Managers

What Is a Context Manager in Python?

A context manager is an object that sets something up and tears it down. It handles the “before” and “after” automatically.

The with statement is how we use a context manager.

with open("file.txt", "r") as file:
    content = file.read()

Here, open("file.txt", "r") is a context manager. It opens the file. Then it gives us the file object. When the block ends, it closes the file. Even if an exception happens.

That is the magic of context managers. They handle cleanup for us.

Why Use Context Managers in Python?

Without context managers, we have to close files manually.

file = open("file.txt", "r")
try:
    content = file.read()
finally:
    file.close()

That works. But it is verbose. And we might forget to close the file.

With a context manager, we just write:

with open("file.txt", "r") as file:
    content = file.read()

Much cleaner. And safer.

How Does the with Statement Work in Python?

The with statement works with any object that has __enter__ and __exit__ methods.

• __enter__ is called at the start of the block. It sets things up.
• __exit__ is called at the end of the block. It cleans things up.

When we write:

with some_object as variable:
    # do something

Python does this:

1. Calls some_object.__enter__() and assigns the result to variable.
2. Runs the block.
3. Calls some_object.__exit__() at the end. Even if an exception happens.

That is how it works.

How to Create a Context Manager Using a Class in Python

We can create our own context manager by defining a class with __enter__ and __exit__.

Let’s create a simple context manager that prints a message before and after.

class MyContext:
    def __enter__(self):
        print("Entering the block")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print("Exiting the block")

with MyContext() as context:
    print("Inside the block")

Entering the block
Inside the block
Exiting the block

That is the basic pattern.

The __exit__ method takes three arguments: the exception type, the exception value, and the traceback. If no exception happened, they are all None.

We can use them to handle exceptions.

class MyContext:
    def __enter__(self):
        print("Entering")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            print(f"An exception occurred: {exc_value}")
        print("Exiting")
        return False   # do not suppress the exception

with MyContext() as context:
    print("Inside")
    raise ValueError("Something went wrong")

Entering
Inside
An exception occurred: Something went wrong
Exiting
Traceback (most recent call last):
  ...
ValueError: Something went wrong

The __exit__ method ran. It printed the exception message. Then it returned False, so the exception was not suppressed. The program still crashed.

If we return True, the exception is suppressed.

class MyContext:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            print(f"Caught: {exc_value}")
        return True   # suppress the exception

with MyContext() as context:
    raise ValueError("Something went wrong")

print("After the block")

Caught: Something went wrong
After the block

The exception was caught and suppressed. The program continued.

Python Context Manager Example – Timing a Block

Let’s create a context manager that measures how long a block takes.

import time

class Timer:
    def __enter__(self):
        self.start = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end = time.time()
        print(f"Elapsed time: {self.end - self.start:.4f} seconds")

with Timer():
    total = 0
    for i in range(1000000):
        total += i

Elapsed time: 0.0452 seconds

That is a useful context manager.

How to Use contextlib for Simpler Context Managers in Python

Python has a module called contextlib. It provides a decorator called @contextmanager that lets us write context managers as functions.

from contextlib import contextmanager

@contextmanager
def my_context():
    print("Entering")
    yield
    print("Exiting")

with my_context():
    print("Inside")

Entering
Inside
Exiting

That is much simpler than writing a class.

We can also handle exceptions.

from contextlib import contextmanager

@contextmanager
def my_context():
    print("Entering")
    try:
        yield
    except Exception as e:
        print(f"Caught: {e}")
    finally:
        print("Exiting")

with my_context():
    print("Inside")
    raise ValueError("Something went wrong")

Entering
Inside
Caught: Something went wrong
Exiting

The exception is caught inside the context manager. But it is not suppressed. So the program would still crash after the block.

If we want to suppress it, we can use return after the except.

@contextmanager
def my_context():
    print("Entering")
    try:
        yield
    except Exception as e:
        print(f"Caught: {e}")
        return
    finally:
        print("Exiting")

Now the exception is suppressed.

Python Context Manager Example – Database Connection

from contextlib import contextmanager

@contextmanager
def database_connection():
    print("Connecting to database")
    connection = "fake_connection"
    try:
        yield connection
    finally:
        print("Closing connection")

with database_connection() as conn:
    print(f"Using {conn}")

Connecting to database
Using fake_connection
Closing connection

That is a clean way to manage resources.

Common Mistakes with Context Managers in Python

file = open("data.txt", "r")
content = file.read()
# forgot to close

with open("data.txt", "r") as file:
    content = file.read()

class MyContext:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # no handling
        pass

@contextmanager
def my_context():
    yield

from contextlib import contextmanager

Python Context Manager Example – Temporary File

Let’s build a context manager that creates a temporary file and deletes it when done.

import os
from contextlib import contextmanager

@contextmanager
def temporary_file(filename):
    with open(filename, "w") as file:
        file.write("Temporary data")
    print(f"Created {filename}")
    try:
        yield filename
    finally:
        os.remove(filename)
        print(f"Deleted {filename}")

with temporary_file("temp.txt") as fname:
    print(f"Using {fname}")
    with open(fname, "r") as file:
        print(file.read())

Created temp.txt
Using temp.txt
Temporary data
Deleted temp.txt

That is a useful context manager. It creates a file, lets us use it, and deletes it when we are done.

Context managers are a powerful tool. They keep our code clean and safe. And with contextlib, they are easy to write.

Next, we will learn Python Modules, so we can split Python code into separate files and reuse it cleanly.