Python Custom Context Managers

In the last lesson, we learned about Python Closures. Now let us create custom context managers for reliable setup and cleanup.

🧩 What Is a Custom Context Manager in Python?

A custom context manager is a class or function that we can use with the with statement. It sets something up at the start and tears it down at the end.

We use it when we want to manage a resource, like a file, a network connection, or a lock.

🎯 How to Create a Context Manager Using a Class

We create a class with __enter__ and __exit__ methods.

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

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

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

Entering
Inside
Exiting

The __enter__ method runs at the start. The __exit__ method runs at the end. Even if an exception happens.

⚙️ Handling Exceptions in exit

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

We can check for exceptions and handle them.

class MyContext:
    def __enter__(self):
        return self

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

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

Caught: Something went wrong
Traceback (most recent call last):
  ...
ValueError: Something went wrong

If we return True, the exception is suppressed.

return True

Now the program continues after the block.

🔧 Using contextlib for Simpler Context Managers

We can use the @contextmanager decorator from contextlib to 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.

We can also handle exceptions.

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

Now if an exception happens, it is caught and printed. Then the finally block runs. The exception is not suppressed unless we return from the except block.

💡 A Real-Life Example – Timer Context Manager

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.time()
    try:
        yield
    finally:
        end = time.time()
        print(f"Elapsed: {end - start:.4f} seconds")

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

Elapsed: 0.0452 seconds

That is a useful context manager for profiling code.

🚀 Another 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 Custom Context Managers

🧪 Mistake 1: Forgetting to return True to suppress exceptions

If we want to suppress an exception, we must return True from __exit__.

🔍 Mistake 2: Not using try/finally in contextmanager

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

If an exception happens inside the with block, the code after yield does not run. That means cleanup does not happen.

Fix: Use try/finally.

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

📋 Mistake 3: Forgetting to import contextmanager

@contextmanager
def my_context():
    yield

This gives an error if we did not import.

Fix:

from contextlib import contextmanager

Next, we will learn about Python Dataclasses and create data-focused classes with less boilerplate.