Python's Context Managers Beyond `with`: Creating Custom Context Managers for Resource Management

You know the feeling when you're working with files, database connections, or network sockets and you just wish there was a cleaner way to handle setup and cleanup? That's exactly what Python's context managers solve. But here's the thing—most developers only scratch the surface by using with open() or similar built-in examples. The real power lies in creating your own.

Let me show you something interesting. At PythonSkillset, we once had a situation where we needed to manage temporary database connections across multiple microservices. The standard with statement helped, but we needed more control. That's when we dove deep into custom context managers.

The Basics We All Know

First, let's quickly recap what you're probably already using:

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

This works because the file object is a context manager. It knows how to open and close itself properly. But what if you need to manage things that don't come with built-in context management? Think database connections, temporary directories, or even just measuring execution time.

Two Ways to Build Your Own

Python gives us two approaches: using a class with __enter__ and __exit__ methods, or using the contextlib module with @contextmanager decorator. Let's look at both, because each has its sweet spot.

The Class-Based Approach

Here's a practical example from a real project at PythonSkillset—a simple database connection manager:

class DatabaseConnection:
    def __init__(self, host, port, database):
        self.host = host
        self.port = port
        self.database = database
        self.connection = None

    def __enter__(self):
        print(f"Connecting to {self.host}:{self.port}/{self.database}")
        self.connection = {"host": self.host, "port": self.port, "connected": True}
        return self.connection

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Closing connection")
        self.connection["connected"] = False
        if exc_type:
            print(f"Error occurred: {exc_val}")
        return False  # Don't suppress exceptions

# Usage
with DatabaseConnection("localhost", 5432, "pythonskillset") as conn:
    print(f"Connection status: {conn['connected']}")

Notice something? The __exit__ method gets three arguments: exception type, value, and traceback. If no exception occurred, they're all None. The return value of False means any exception will be propagated—usually what you want.

The Generator-Based Approach (My Favorite)

Using contextmanager from contextlib is often cleaner for simpler cases:

from contextlib import contextmanager
import time

@contextmanager
def timer(name="operation"):
    start = time.perf_counter()
    print(f"Starting {name}")
    try:
        yield
    finally:
        duration = time.perf_counter() - start
        print(f"{name} took {duration:.2f} seconds")

# Usage
with timer("data_processing"):
    data = [i for i in range(1000000)]
    total = sum(data)
    print(f"Sum calculated: {total}")

The key is the yield statement—it's where your code runs. Everything before yield is setup, everything after is cleanup. And it always runs, even if an exception happens.

Advanced Techniques Worth Knowing

Resource Pooling

At PythonSkillset, we built a connection pool for our database queries. Custom context managers made this elegant:

class ConnectionPool:
    def __init__(self, max_connections=5):
        self.max_connections = max_connections
        self.available = []
        self.in_use = set()

    def acquire(self):
        if self.available:
            conn = self.available.pop()
        else:
            conn = {"id": len(self.in_use) + len(self.available) + 1}
        self.in_use.add(conn["id"])
        return conn

    def release(self, conn):
        self.in_use.discard(conn["id"])
        self.available.append(conn)

    @contextmanager
    def get_connection(self):
        conn = self.acquire()
        try:
            yield conn
        finally:
            self.release(conn)

pool = ConnectionPool()
with pool.get_connection() as conn:
    print(f"Using connection #{conn['id']}")

Nested Context Managers

You can stack them like blocks:

@contextmanager
def temporary_change(obj, attr, value):
    old_value = getattr(obj, attr, None)
    setattr(obj, attr, value)
    try:
        yield
    finally:
        setattr(obj, attr, old_value)

class Config:
    debug = False

config = Config()
with temporary_change(config, "debug", True):
    print(f"Debug mode: {config.debug}")
print(f"Back to: {config.debug}")

Real-World Example from PythonSkillset

We once had a scenario where we needed to safely handle file uploads in a web application. Each upload needed a temporary directory that would be automatically cleaned up:

import tempfile
import shutil
from pathlib import Path

@contextmanager
def temporary_workspace(prefix="upload_"):
    temp_dir = tempfile.mkdtemp(prefix=prefix)
    print(f"Created workspace: {temp_dir}")
    try:
        yield Path(temp_dir)
    finally:
        print(f"Cleaning up: {temp_dir}")
        shutil.rmtree(temp_dir, ignore_errors=True)

with temporary_workspace() as workspace:
    file_path = workspace / "temp_data.txt"
    file_path.write_text("This will be cleaned up")
    print(f"File written to: {file_path}")
    # Process the file...

Common Pitfalls to Avoid

1. Forgetting to return True for exception handling—If you want to suppress an exception, return True from __exit__. Otherwise, let it propagate.

2. Not handling all exceptions—Your cleanup code should work regardless of whether an exception occurred. That's why finally blocks are crucial.

3. Overcomplicating simple cases—If you're just timing code or logging, a decorated function is cleaner than a class.

When to Use Which?

Use the class-based approach when: - You need to manage state across multiple context entries - You're working with resources that need complex initialization - You want to expose additional methods for the resource

Use the @contextmanager decorator when: - Your setup/cleanup is straightforward - You're wrapping a simple process - You want cleaner, more readable code

The Bottom Line

Custom context managers aren't just about closing files or connections—they're about making your code more reliable and readable. At PythonSkillset, we've found they reduce bugs by ensuring cleanup always happens, even when things go wrong. Plus, they make your code more Pythonic.

Next time you find yourself writing try...finally blocks, stop and ask: could this be a context manager? Chances are, the answer is yes, and your future self will thank you for the cleaner code.

Start small. Wrap one of your existing cleanup patterns. See how it feels. I bet you'll find yourself reaching for context managers more than you expected.