Python - Context Manager (with statement)

Key Insights

• Context managers automate resource setup and teardown using the with statement, guaranteeing cleanup even when exceptions occur • Implement custom context managers via __enter__ and __exit__ methods or the @contextmanager decorator for cleaner resource management • Context managers eliminate common bugs like unclosed files, unreleased locks, and uncommitted database transactions through deterministic cleanup

Understanding Context Managers

Context managers provide a protocol for wrapping code execution with setup and teardown logic. The with statement invokes __enter__ before the block executes and guarantees __exit__ runs afterward, regardless of exceptions.

# Without context manager - manual cleanup
file = open('data.txt', 'r')
try:
    content = file.read()
    process(content)
finally:
    file.close()

# With context manager - automatic cleanup
with open('data.txt', 'r') as file:
    content = file.read()
    process(content)
# file.close() called automatically

The context manager protocol ensures resources are released even if exceptions occur within the block. This prevents resource leaks and makes code more robust.

The Context Manager Protocol

Implement __enter__ and __exit__ to create a context manager class. The __enter__ method runs before the block, __exit__ runs after.

class DatabaseConnection:
    def __init__(self, connection_string):
        self.connection_string = connection_string
        self.connection = None
    
    def __enter__(self):
        # Setup: acquire resource
        self.connection = self._connect(self.connection_string)
        print("Connection established")
        return self.connection  # Value assigned to 'as' variable
    
    def __exit__(self, exc_type, exc_value, traceback):
        # Teardown: release resource
        if self.connection:
            self.connection.close()
            print("Connection closed")
        
        # Return False to propagate exceptions
        # Return True to suppress exceptions
        return False
    
    def _connect(self, connection_string):
        # Simulate connection
        return {"connected": True, "conn_str": connection_string}

# Usage
with DatabaseConnection("postgresql://localhost/mydb") as conn:
    print(f"Using connection: {conn}")
    # Perform database operations
# Connection automatically closed

The __exit__ method receives exception information if an error occurs. Return True to suppress the exception, False to propagate it.

Using contextlib.contextmanager

The @contextmanager decorator converts generator functions into context managers, eliminating boilerplate code.

from contextlib import contextmanager
import time

@contextmanager
def timer(label):
    start = time.time()
    try:
        yield  # Code block executes here
    finally:
        end = time.time()
        print(f"{label}: {end - start:.4f} seconds")

# Usage
with timer("Data processing"):
    # Simulate work
    time.sleep(0.5)
    result = sum(range(1000000))
# Output: Data processing: 0.5XXX seconds

The yield statement separates setup from teardown. Code before yield runs in __enter__, code after runs in __exit__. The finally block ensures cleanup executes even with exceptions.

Practical Example: Transaction Manager

Context managers excel at managing database transactions with automatic rollback on errors.

from contextlib import contextmanager
import sqlite3

class Database:
    def __init__(self, db_path):
        self.db_path = db_path
        self.connection = sqlite3.connect(db_path)
    
    @contextmanager
    def transaction(self):
        cursor = self.connection.cursor()
        try:
            yield cursor
            self.connection.commit()
            print("Transaction committed")
        except Exception as e:
            self.connection.rollback()
            print(f"Transaction rolled back: {e}")
            raise
        finally:
            cursor.close()

# Setup database
db = Database(':memory:')
with db.connection:
    db.connection.execute('''
        CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, balance REAL)
    ''')

# Successful transaction
with db.transaction() as cursor:
    cursor.execute("INSERT INTO users (name, balance) VALUES (?, ?)", 
                   ("Alice", 100.0))
    cursor.execute("INSERT INTO users (name, balance) VALUES (?, ?)", 
                   ("Bob", 50.0))
# Output: Transaction committed

# Failed transaction with rollback
try:
    with db.transaction() as cursor:
        cursor.execute("UPDATE users SET balance = balance - 30 WHERE name = ?", 
                       ("Alice",))
        # Simulate error
        raise ValueError("Insufficient funds")
        cursor.execute("UPDATE users SET balance = balance + 30 WHERE name = ?", 
                       ("Bob",))
except ValueError:
    pass
# Output: Transaction rolled back: Insufficient funds

Multiple Context Managers

Stack multiple context managers or use parentheses for complex resource management.

# Stacked context managers
with open('input.txt', 'r') as infile:
    with open('output.txt', 'w') as outfile:
        for line in infile:
            outfile.write(line.upper())

# Python 3.1+ syntax
with open('input.txt', 'r') as infile, \
     open('output.txt', 'w') as outfile:
    for line in infile:
        outfile.write(line.upper())

# Python 3.10+ parenthesized syntax
with (
    open('input.txt', 'r') as infile,
    open('output.txt', 'w') as outfile,
):
    for line in infile:
        outfile.write(line.upper())

Advanced: Reentrant and Reusable Context Managers

Most context managers aren’t reentrant (can’t be nested) or reusable (can’t be used multiple times). Use contextlib utilities for specialized behavior.

from contextlib import contextmanager
from threading import Lock

# Non-reentrant lock
lock = Lock()

@contextmanager
def acquire_lock():
    lock.acquire()
    try:
        yield
    finally:
        lock.release()

# This would deadlock - non-reentrant
# with acquire_lock():
#     with acquire_lock():  # Deadlock!
#         pass

# Reentrant lock using RLock
from threading import RLock

class ReentrantLockManager:
    def __init__(self):
        self.lock = RLock()
    
    def __enter__(self):
        self.lock.acquire()
        return self
    
    def __exit__(self, *args):
        self.lock.release()
        return False

# This works - reentrant
rlock = ReentrantLockManager()
with rlock:
    print("Outer lock acquired")
    with rlock:
        print("Inner lock acquired (same thread)")
    print("Inner lock released")
print("Outer lock released")

Context Managers for Testing

Context managers simplify test setup and teardown, ensuring clean state between tests.

from contextlib import contextmanager
import tempfile
import shutil
import os

@contextmanager
def temporary_directory():
    """Create a temporary directory, yield it, then clean up."""
    temp_dir = tempfile.mkdtemp()
    try:
        yield temp_dir
    finally:
        shutil.rmtree(temp_dir)

@contextmanager
def change_directory(path):
    """Change to a directory temporarily, then restore."""
    original = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(original)

# Usage in tests
def test_file_operations():
    with temporary_directory() as temp_dir:
        with change_directory(temp_dir):
            # Create test files
            with open('test.txt', 'w') as f:
                f.write('test data')
            
            # Perform operations
            assert os.path.exists('test.txt')
            
        # Directory and files automatically cleaned up
        assert not os.path.exists(temp_dir)

Error Handling and Suppression

Control exception propagation through __exit__ return values or use contextlib.suppress.

from contextlib import contextmanager, suppress

@contextmanager
def ignore_errors(*exceptions):
    """Suppress specific exceptions."""
    try:
        yield
    except exceptions:
        pass  # Suppress by not re-raising

# Usage
with ignore_errors(FileNotFoundError, PermissionError):
    os.remove('nonexistent_file.txt')
# No exception raised

# Using contextlib.suppress (cleaner)
with suppress(FileNotFoundError, PermissionError):
    os.remove('nonexistent_file.txt')

# Conditional suppression
@contextmanager
def conditional_suppress(condition, *exceptions):
    try:
        yield
    except exceptions as e:
        if not condition:
            raise
        # Suppress if condition is True

with conditional_suppress(True, ValueError):
    raise ValueError("This is suppressed")

try:
    with conditional_suppress(False, ValueError):
        raise ValueError("This propagates")
except ValueError:
    print("Exception not suppressed")

Performance Considerations

Context managers add minimal overhead. The protocol involves two method calls (__enter__ and __exit__), which is negligible compared to resource acquisition costs.

import time
from contextlib import contextmanager

@contextmanager
def measure_overhead():
    start = time.perf_counter()
    yield
    overhead = time.perf_counter() - start
    return overhead

# Measure context manager overhead
iterations = 1000000

start = time.perf_counter()
for _ in range(iterations):
    with suppress():
        pass
context_time = time.perf_counter() - start

start = time.perf_counter()
for _ in range(iterations):
    try:
        pass
    except:
        pass
try_time = time.perf_counter() - start

print(f"Context manager: {context_time:.4f}s")
print(f"Try/except: {try_time:.4f}s")
print(f"Overhead per iteration: {(context_time - try_time) / iterations * 1e9:.2f}ns")

Context managers provide deterministic resource management with negligible performance cost. Use them for files, locks, database connections, network sockets, and any resource requiring guaranteed cleanup. The with statement makes resource management explicit, preventing leaks and improving code reliability.