contextlib: Create Your Own Context Manager

The with statement is one of Python’s most expressive features. It guarantees that setup and cleanup logic runs correctly, even when exceptions occur.

Most developers use it with files or database connections. Under the hood, context managers replace repetitive try/finally blocks with a cleaner abstraction.

With contextlib, you can build your own context managers and apply the same pattern to any problem.

Why Context Managers Exist

Without a context manager, you might write:

resource = acquire()
try:
    process(resource)
finally:
    release(resource)

This works, but it is repetitive and easy to forget.

The with statement formalizes this pattern. A context manager defines what happens when entering and exiting a block of code.

The Traditional Approach: enter and exit

You can create a context manager by defining two special methods:

class Timer:
    def __enter__(self):
        print("Starting")
        return self

    def __exit__(self, exc_type, exc, tb):
        print("Finished")

This is powerful but verbose.

For simpler use cases, contextlib provides a cleaner solution.

Using @contextmanager

The @contextmanager decorator lets you define setup and cleanup logic around a single yield.

Code before yield runs when entering the block.
Code after yield runs when exiting the block, even if an exception occurs.

Example:

import time
import tracemalloc
from contextlib import contextmanager

@contextmanager
def performance_tracker(label="Block"):
    start_time = time.perf_counter()
    tracemalloc.start()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start_time
        current, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        print(f"{label} took {elapsed:.4f}s | Peak memory: {peak / 1024:.1f} KB")

Usage:

with performance_tracker("Data Processing"):
    data = [x ** 2 for x in range(1_000_000)]

The tracked code runs normally. When the block finishes, cleanup and reporting happen automatically.

Why This Is Powerful

Custom context managers allow you to:

• Track execution time consistently
• Measure memory usage
• Wrap logging or metrics
• Manage temporary state safely
• Encapsulate cross-cutting concerns cleanly

In data-heavy workflows, this pattern helps monitor transformations without modifying core logic.

I built a small performance measurement context in one of my projects and ended up using it everywhere. It is one of those tools that pays off immediately.

Best Practices

• Always place cleanup logic inside finally
• Keep context managers focused on one responsibility
• Avoid embedding complex business logic inside them
• Use descriptive labels for observability

Final Take

contextlib extends the power of the with statement beyond files and connections.

By encapsulating setup and teardown logic in a reusable abstraction, you reduce repetition and improve reliability.

Context managers are not just syntactic convenience. They are a disciplined way to manage resources and cross-cutting behavior in Python.