DEV Community

Cover image for The Custom Portal Design: Building Your Own Context Managers
Aaron Rose
Aaron Rose

Posted on

The Custom Portal Design: Building Your Own Context Managers

Timothy had mastered using Python's built-in context managers, but the head librarian's next request stumped him. "We need to time how long our cataloging operations take, but only in production mode. The timing code clutters every function. Can you make it cleaner?"

Margaret led him to a workshop labeled "The Custom Portal Design," where librarians crafted specialized self-closing chambers for unique library needs. "You can build your own context managers," she explained. "Design portals that handle any setup and teardown pattern you need."

The Timing Problem

Timothy's timing code was repetitive and error-prone:

import time

# Note: PRODUCTION_MODE and log_performance() are placeholder configuration
# In practice, replace with your actual settings and logging functions

def catalog_book(title, author):
    start_time = time.time()
    try:
        # Actual cataloging work
        result = add_to_database(title, author)
        verify_entry(result)
        update_index(result)
        return result
    finally:
        elapsed = time.time() - start_time
        if PRODUCTION_MODE:
            log_performance("catalog_book", elapsed)

def search_catalog(query):
    start_time = time.time()
    try:
        # Search logic
        results = query_database(query)
        return results
    finally:
        elapsed = time.time() - start_time
        if PRODUCTION_MODE:
            log_performance("search_catalog", elapsed)
Enter fullscreen mode Exit fullscreen mode

Every function duplicated the timing logic. Timothy wanted to write simply:

def catalog_book(title, author):
    with Timer("catalog_book"):
        result = add_to_database(title, author)
        verify_entry(result)
        update_index(result)
        return result
Enter fullscreen mode Exit fullscreen mode

Margaret showed him how to build the Timer context manager.

Building a Class-Based Context Manager

"Context managers," Margaret explained, "are just classes implementing __enter__ and __exit__ methods."

import time

class Timer:
    def __init__(self, operation_name):
        self.operation_name = operation_name
        self.start_time = None

    def __enter__(self):
        self.start_time = time.time()
        return self  # Optional: return self for access inside with block

    def __exit__(self, exc_type, exc_value, traceback):
        elapsed = time.time() - self.start_time
        if PRODUCTION_MODE:
            log_performance(self.operation_name, elapsed)
        return False  # Don't suppress exceptions

# Use it
with Timer("catalog_book"):
    add_to_database(title, author)
    verify_entry(result)
Enter fullscreen mode Exit fullscreen mode

Timothy traced the execution:

  1. Timer("catalog_book") created an instance
  2. __enter__ recorded the start time when entering the block
  3. The cataloging code executed
  4. __exit__ calculated elapsed time and logged it when leaving
  5. Exceptions propagated normally (return False)

"The __init__ method," Margaret noted, "receives configuration. The __enter__ method does setup. The __exit__ method does teardown and cleanup."

The Contextlib Decorator Approach

Margaret revealed a cleaner way to build context managers for simple cases:

from contextlib import contextmanager
import time

@contextmanager
def timer(operation_name):
    start_time = time.time()
    try:
        yield  # This is where the with block executes
    finally:
        elapsed = time.time() - start_time
        if PRODUCTION_MODE:
            log_performance(operation_name, elapsed)

# Use it identically
with timer("catalog_book"):
    add_to_database(title, author)
Enter fullscreen mode Exit fullscreen mode

The decorator transformed a generator function into a context manager:

  • Code before yield became the __enter__ method
  • Code after yield became the __exit__ method
  • The finally block ensured cleanup happened even with exceptions

"When you don't need a class with state," Margaret explained, "the decorator approach is more concise."

Returning Values from Context Managers

Timothy discovered the yield statement could provide values to the with block:

from contextlib import contextmanager
import time

@contextmanager
def timer(operation_name):
    timing_data = {"start": time.time()}

    yield timing_data  # Provide access to timing info

    timing_data["end"] = time.time()
    timing_data["elapsed"] = timing_data["end"] - timing_data["start"]
    if PRODUCTION_MODE:
        log_performance(operation_name, timing_data["elapsed"])

# Access the timing data
with timer("catalog_book") as timing:
    add_to_database(title, author)
    print(f"Started at {timing['start']}")
Enter fullscreen mode Exit fullscreen mode

Whatever the generator yielded became available via the as clause. In class-based context managers, __enter__'s return value served the same purpose.

The Temporary Settings Pattern

Margaret showed Timothy a practical pattern—temporarily changing configuration:

from contextlib import contextmanager

# Assumes get_setting() and set_setting() functions exist in your codebase

@contextmanager
def temporary_setting(setting_name, temporary_value):
    # Save original value
    original_value = get_setting(setting_name)

    # Set temporary value
    set_setting(setting_name, temporary_value)

    try:
        yield
    finally:
        # Restore original value
        set_setting(setting_name, original_value)

# Use it
with temporary_setting("DEBUG_MODE", True):
    # Debug mode is on
    diagnose_catalog_issue()
# Debug mode restored to original state
Enter fullscreen mode Exit fullscreen mode

The pattern was setup, execute, restore—perfect for context managers.

The Database Transaction Manager

Timothy built a context manager for database transactions:

class DatabaseTransaction:
    def __init__(self, connection):
        self.connection = connection

    def __enter__(self):
        self.connection.begin_transaction()
        return self.connection

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            # No exception - commit
            self.connection.commit()
        else:
            # Exception occurred - rollback
            self.connection.rollback()
        return False  # Let exception propagate after rollback

# Use it
with DatabaseTransaction(db_connection) as db:
    db.execute("INSERT INTO books VALUES (?, ?)", (title, author))
    db.execute("UPDATE catalog SET count = count + 1")
    # Automatically commits if successful, rolls back if exception occurs
Enter fullscreen mode Exit fullscreen mode

The context manager examined exc_type to decide whether to commit or rollback. Transactions became automatic and safe.

Exception Handling and Suppression

Margaret showed Timothy how to handle specific exceptions:

from contextlib import contextmanager

@contextmanager
def suppress_file_not_found():
    try:
        yield
    except FileNotFoundError as e:
        # Log it but don't crash
        print(f"File not found (expected): {e}")
        # By not re-raising, we suppress the exception

with suppress_file_not_found():
    process_optional_config_file()
# Program continues even if file doesn't exist
Enter fullscreen mode Exit fullscreen mode

Timothy learned that Python's standard library included contextlib.suppress for this pattern:

from contextlib import suppress

with suppress(FileNotFoundError, PermissionError):
    delete_temporary_file()
# Silently ignores these exceptions
Enter fullscreen mode Exit fullscreen mode

"But remember," Margaret cautioned, "suppression should be rare and intentional. Most context managers let exceptions propagate."

The Lock Manager Pattern

Timothy created a context manager for file locking:

import fcntl

# Note: fcntl is Unix/Linux specific. For cross-platform file locking,
# consider using the 'filelock' or 'portalocker' package from PyPI.

class FileLock:
    def __init__(self, filename):
        self.filename = filename
        self.file_handle = None

    def __enter__(self):
        self.file_handle = open(self.filename, 'a')
        fcntl.flock(self.file_handle.fileno(), fcntl.LOCK_EX)
        return self.file_handle

    def __exit__(self, exc_type, exc_value, traceback):
        if self.file_handle:
            fcntl.flock(self.file_handle.fileno(), fcntl.LOCK_UN)
            self.file_handle.close()
        return False

# Ensure only one process modifies the catalog at a time
with FileLock("catalog.lock"):
    update_shared_catalog()
# Lock released, file closed
Enter fullscreen mode Exit fullscreen mode

The context manager acquired an exclusive lock on entry and released it on exit, preventing concurrent modification issues.

The Nested Context Manager

Margaret demonstrated that context managers could use other context managers internally:

import os
from contextlib import contextmanager

@contextmanager
def atomic_catalog_update(catalog_file):
    # Note: Assumes catalog_file already exists
    # For creating new files atomically, additional logic would be needed

    with open(catalog_file, 'r') as f:
        original_content = f.read()

    temp_file = catalog_file + ".tmp"

    try:
        with open(temp_file, 'w') as f:
            yield f

        # Success - replace original
        os.replace(temp_file, catalog_file)
    except Exception:
        # Failure - restore original
        with open(catalog_file, 'w') as f:
            f.write(original_content)
        raise
    finally:
        # Clean up temp file if it exists
        if os.path.exists(temp_file):
            os.remove(temp_file)

with atomic_catalog_update("catalog.txt") as catalog:
    catalog.write("Updated catalog data")
# Changes are atomic - either fully applied or fully reverted
Enter fullscreen mode Exit fullscreen mode

Context managers could compose, building complex resource management from simpler pieces.

The Reusable vs Single-Use Pattern

Timothy learned the difference between reusable and single-use context managers:

# Reusable - can use multiple times
timer = Timer("operation")

with timer:
    do_work_1()

with timer:
    do_work_2()

# Single-use - works only once
@contextmanager
def single_use():
    yield
    print("Can only use once")

manager = single_use()
with manager:
    do_work()

with manager:  # Error! Generator already exhausted
    do_more_work()
Enter fullscreen mode Exit fullscreen mode

Class-based context managers were naturally reusable. Function-based ones (with @contextmanager) were single-use unless explicitly designed otherwise.

The Resource Pool Pattern

Margaret showed Timothy an advanced pattern—managing a pool of resources:

from contextlib import contextmanager

# Note: create_connection() would be your actual connection creation function
# e.g., for databases: sqlite3.connect(), psycopg2.connect(), etc.

class ConnectionPool:
    def __init__(self, max_connections):
        self.pool = [create_connection() for _ in range(max_connections)]
        self.available = self.pool.copy()

    @contextmanager
    def connection(self):
        if not self.available:
            raise Exception("No connections available")

        conn = self.available.pop()
        try:
            yield conn
        finally:
            self.available.append(conn)

pool = ConnectionPool(max_connections=5)

# Borrow and return connections automatically
with pool.connection() as conn:
    query_database(conn)
# Connection returned to pool
Enter fullscreen mode Exit fullscreen mode

The context manager borrowed a connection from the pool and guaranteed its return, even if exceptions occurred.

Timothy's Custom Context Manager Wisdom

Through mastering the Custom Portal Design, Timothy learned essential principles:

Two ways to build context managers: Class-based (__enter__/__exit__) or decorator-based (@contextmanager).

Use decorators for simple cases: When you don't need reusable state, @contextmanager is cleaner.

Use classes for complex state: When managing multiple resources or needing reusability, use classes.

The yield point matters: Code before yield is setup, code after is teardown.

Always use try/finally: In @contextmanager functions, wrap the yield to ensure cleanup.

Return False by default: Let exceptions propagate unless you have a specific reason to suppress.

Examine exc_type for conditional cleanup: Commit on success, rollback on failure.

Context managers compose: Build complex resource management from simpler context managers.

Consider reusability: Class-based managers work multiple times; generator-based ones are typically single-use.

Timothy's exploration of custom context managers revealed that any setup-and-teardown pattern could become a clean, reusable context manager. The Custom Portal Design transformed repetitive resource management code into declarative, self-documenting patterns. Whether timing operations, managing transactions, or controlling access to shared resources, context managers guaranteed proper cleanup while keeping code focused on business logic rather than bookkeeping.


Aaron Rose is a software engineer and technology writer at tech-reader.blog and the author of Think Like a Genius.

Top comments (0)