Python
  • Understand the purpose and benefits of context managers
  • Create custom context managers using classes and decorators
  • Use context managers for resource management
  • Apply context managers in real-world scenarios

Context Managers

Introduction

Imagine you're a librarian managing a busy library. Every time someone borrows a book, you need to record who took it, when they took it, and ensure they return it properly. If they don't return it on time, you need to send reminders. Managing all these details manually would be overwhelming.

Context managers in Python are like that librarian - they handle the setup and cleanup of resources automatically. They ensure that resources are properly acquired and released, even if errors occur. This is one of Python's most powerful features for writing robust, clean code.

The Problem Context Managers Solve

Before context managers, resource management often looked like this:

file = open('data.txt', 'r')
try:
    data = file.read()
    # Process data
    print(data)
finally:
    file.close()  # Must remember to close.

What if an exception occurs? What if there are multiple resources? What if the cleanup logic is complex? Context managers solve these problems elegantly.

The with Statement

The with statement creates a context that automatically handles setup and cleanup:

with open('data.txt', 'r') as file:
    data = file.read()
    print(data)
# File is automatically closed here

This is equivalent to:

file = open('data.txt', 'r')  # Setup
try:
    data = file.read()
    print(data)
finally:
    file.close()  # Cleanup

But much cleaner and safer.

Creating Context Managers

Class-Based Context Managers

To create a context manager, implement __enter__ and __exit__ methods:

class DatabaseConnection:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.connection = None
    
    def __enter__(self):
        """Called when entering the with block."""
        print(f"Connecting to {self.host}:{self.port}")
        # Simulate connection
        self.connection = f"Connection to {self.host}:{self.port}"
        return self.connection  # This becomes the 'as' variable
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Called when exiting the with block."""
        print(f"Disconnecting from {self.host}:{self.port}")
        self.connection = None
        # Return False to propagate exceptions, True to suppress them
        return False

# Usage
with DatabaseConnection("localhost", 5432) as conn:
    print(f"Using {conn}")
    # Do database operations
# Connection automatically closed

Decorator-Based Context Managers

The contextlib.contextmanager decorator lets you create context managers from generator functions:

from contextlib import contextmanager

@contextmanager
def database_connection(host, port):
    """Context manager using decorator approach."""
    print(f"Connecting to {host}:{port}")
    connection = f"Connection to {host}:{port}"
    try:
        yield connection  # Code before yield = setup, after = cleanup
    finally:
        print(f"Disconnecting from {host}:{port}")

# Usage
with database_connection("localhost", 5432) as conn:
    print(f"Using {conn}")
    # Do database operations

Exception Handling in Context Managers

Context managers can handle exceptions that occur within the with block:

class SafeFileHandler:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.file = None
    
    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            self.file.close()
        
        # Log any exceptions
        if exc_type:
            print(f"Exception occurred: {exc_type.__name__}: {exc_val}")
            # Return False to let exception propagate
            return False
        
        print("File operation completed successfully")
        return True

# Usage
with SafeFileHandler('data.txt', 'w') as f:
    f.write("Hello World")
    # If an exception occurs here, __exit__ will still close the file

Real-World Applications

Database Transactions

class DatabaseTransaction:
    def __init__(self, db_connection):
        self.db = db_connection
        self.transaction_started = False
    
    def __enter__(self):
        self.db.execute("BEGIN TRANSACTION")
        self.transaction_started = True
        return self.db
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            # Rollback on error
            self.db.execute("ROLLBACK")
            print("Transaction rolled back due to error")
        else:
            # Commit on success
            self.db.execute("COMMIT")
            print("Transaction committed successfully")
        
        self.transaction_started = False

# Usage
with DatabaseTransaction(db_conn) as db:
    db.execute("INSERT INTO users (name) VALUES (?)", ("Alice",))
    db.execute("INSERT INTO users (name) VALUES (?)", ("Bob",))
    # If any operation fails, entire transaction rolls back

Temporary File Management

import tempfile
import os

@contextmanager
def temporary_file(suffix='', prefix='tmp'):
    """Create a temporary file that gets deleted automatically."""
    fd, path = tempfile.mkstemp(suffix=suffix, prefix=prefix)
    try:
        # Close the file descriptor so it can be opened normally
        os.close(fd)
        yield path
    finally:
        # Clean up the temporary file
        if os.path.exists(path):
            os.unlink(path)

# Usage
with temporary_file(suffix='.txt') as temp_path:
    with open(temp_path, 'w') as f:
        f.write("Temporary data")
    
    # File exists here
    print(f"Created temporary file: {temp_path}")

# File is automatically deleted here

Network Connection Pool

class ConnectionPool:
    def __init__(self, max_connections=10):
        self.max_connections = max_connections
        self.available = []
        self.in_use = set()
    
    def __enter__(self):
        if len(self.available) > 0:
            conn = self.available.pop()
        elif len(self.in_use) < self.max_connections:
            conn = self._create_connection()
        else:
            raise RuntimeError("No connections available")
        
        self.in_use.add(conn)
        return conn
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # Return connection to pool
        conn = next(iter(self.in_use))  # Get one connection
        self.in_use.remove(conn)
        self.available.append(conn)

# Usage
pool = ConnectionPool(max_connections=5)
with pool as conn1:
    # Use connection
    pass
# Connection returned to pool automatically

Built-in Context Managers

Python provides several built-in context managers:

File Operations

# Files are context managers
with open('file.txt', 'r') as f:
    content = f.read()

# Equivalent to:
f = open('file.txt', 'r')
try:
    content = f.read()
finally:
    f.close()

Thread Locks

import threading

lock = threading.Lock()

with lock:
    # Critical section - only one thread at a time
    shared_data += 1

Decimal Context

from decimal import Decimal, getcontext, localcontext

with localcontext() as ctx:
    ctx.prec = 42  # Temporary precision change
    result = Decimal('1') / Decimal('7')
# Precision returns to previous value

Context Manager Utilities

contextlib Utilities

from contextlib import contextmanager, suppress, redirect_stdout
import io

# Suppress specific exceptions
with suppress(FileNotFoundError):
    with open('nonexistent.txt') as f:
        content = f.read()

# Redirect output
output = io.StringIO()
with redirect_stdout(output):
    print("This goes to output variable")
    
print(f"Captured: {output.getvalue()}")

# Nested contexts
@contextmanager
def tag(name):
    print(f"<{name}>")
    yield
    print(f"</{name}>")

with tag("html"):
    with tag("body"):
        print("Content")
# Output: <html><body>Content</body></html>

Multiple Context Managers

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

# Or using backslash continuation
with open('file1.txt') as f1, \
     open('file2.txt') as f2, \
     open('file3.txt') as f3:
    # Work with all three files
    pass

Best Practices

1. Always Use Context Managers for Resources

# Good
with open('file.txt') as f:
    data = f.read()

# Avoid
f = open('file.txt')
data = f.read()
f.close()  # Easy to forget.

2. Handle Exceptions Properly

def __exit__(self, exc_type, exc_val, exc_tb):
    # Always perform cleanup
    self.cleanup()
    
    # Decide whether to suppress exception
    if isinstance(exc_val, ExpectedError):
        return True  # Suppress
    return False    # Propagate

3. Make Context Managers Reusable

@contextmanager
def database_session(url):
    session = create_session(url)
    try:
        yield session
    finally:
        session.close()

# Can be reused with different URLs
with database_session("sqlite:///local.db") as session:
    # Local operations
    pass

with database_session("postgresql://remote/db") as session:
    # Remote operations
    pass

Key Points to Remember

  • Context managers ensure proper resource management with automatic cleanup
  • Use __enter__ for setup and __exit__ for cleanup
  • The contextlib.contextmanager decorator simplifies context manager creation
  • Context managers handle exceptions gracefully
  • Built-in types like files and locks are context managers
  • Always prefer with statements over manual resource management

Context managers handle resource management beautifully, but what about processing large datasets without loading everything into memory? In the next lesson, we'll explore generators - Python's solution for lazy evaluation and memory-efficient programming.