deep-dive2023-07-01·9·348/348

Python context manager (with문) 활용

A comprehensive guide to Python context managers — from basic file handling to custom context managers, contextlib patterns, and error handling.

Python context manager (with문) 활용

Introduction

Context managers are Python's answer to the resource acquisition is initialization (RAII) pattern. They guarantee that setup code runs before a block of code and cleanup code runs after — even if an exception occurs. The with statement is used daily by most Python developers for file handling, database connections, and locks, but many do not realize how powerful and flexible context managers can be.

I started building custom context managers seriously after a Lisbon project where database connections were leaking. The application was opening connections in try blocks without proper cleanup, and under load the connection pool was exhausted. Context managers solved this problem elegantly.

Environment

Python 3.12.3

Problem

Problem 1: Manual resource cleanup is error-prone

# BAD: Manual cleanup
def process_file():
    f = open("data.txt", "r")
    try:
        content = f.read()
        # What if this throws an exception?
        process(content)
    finally:
        f.close()  # This might not run if close() itself fails

Problem 2: Database connection leaks

# BAD: Connection might not be closed on error
def get_user(user_id):
    conn = psycopg2.connect("postgresql://localhost/mydb")
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM users WHERE id = %s", (user_id,))
    result = cursor.fetchone()
    cursor.close()
    conn.close()
    return result

Problem 3: Lock acquisition without release

import threading

lock = threading.Lock()

def critical_section():
    lock.acquire()
    # What if this throws?
    do_something()
    lock.release()  # Never reached on error!

Analysis

A context manager is any object that implements __enter__ and __exit__ methods.

The protocol:

class MyContextManager:
    def __enter__(self):
        """Setup code runs here. Returns value used in 'as' clause."""
        print("Entering context")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Cleanup code runs here. Return True to suppress exceptions."""
        print("Exiting context")
        return False  # Don't suppress exceptions

How with statement works:

with MyContextManager() as cm:
    # __enter__ has been called
    do_something()
# __exit__ is called here, even if an exception occurred

The __exit__ parameters:

  • exc_type: The exception class (e.g., ValueError)
  • exc_val: The exception instance
  • exc_tb: The traceback object
  • If all are None, no exception occurred

Solution

Fix 1: Basic context manager with classes

class FileManager:
    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()
        return False  # Don't suppress exceptions

# Usage
with FileManager("data.txt", "w") as f:
    f.write("Hello, world!")
# File is automatically closed, even if write() raises

Fix 2: Using contextlib.contextmanager decorator

from contextlib import contextmanager

@contextmanager
def managed_resource(resource_name):
    """Context manager using generator syntax."""
    print(f"Acquiring {resource_name}")
    resource = acquire_resource(resource_name)
    try:
        yield resource  # This is the 'as' value
    except Exception as e:
        print(f"Error: {e}")
        raise  # Re-raise after cleanup
    finally:
        print(f"Releasing {resource_name}")
        release_resource(resource)

# Usage
with managed_resource("database") as db:
    db.query("SELECT * FROM users")

Fix 3: Database connection context manager

from contextlib import contextmanager
import psycopg2

@contextmanager
def get_db_connection():
    """Provide a transactional database connection."""
    conn = psycopg2.connect("postgresql://localhost/mydb")
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()

# Usage
with get_db_connection() as conn:
    cursor = conn.cursor()
    cursor.execute("INSERT INTO users (name) VALUES (%s)", ("Alice",))
# Connection is closed, transaction committed or rolled back

Fix 4: Timing context manager

from contextlib import contextmanager
import time

@contextmanager
def timer(label="Block"):
    """Measure execution time of a code block."""
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label} took {elapsed:.4f} seconds")

# Usage
with timer("Data processing"):
    process_large_dataset()
# Prints: Data processing took 2.3456 seconds

Fix 5: Nested context managers

from contextlib import contextmanager

@contextmanager
def database():
    conn = create_connection()
    try:
        yield conn
    finally:
        conn.close()

@contextmanager
def transaction(conn):
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise

# Nested usage
with database() as conn:
    with transaction(conn) as tx:
        tx.execute("INSERT INTO users VALUES (1, 'Alice')")
    # Transaction committed, connection still open
# Connection closed

# Or using contextlib.ExitStack
from contextlib import ExitStack

with ExitStack() as stack:
    conn = stack.enter_context(database())
    tx = stack.enter_context(transaction(conn))
    tx.execute("INSERT INTO users VALUES (1, 'Alice')")

Fix 6: Suppress exceptions with contextlib.suppress

from contextlib import suppress

# Suppress FileNotFoundError
with suppress(FileNotFoundError):
    os.remove("maybe_exists.txt")
# If file doesn't exist, no error — silently ignored

# Equivalent to try/except with pass
try:
    os.remove("maybe_exists.txt")
except FileNotFoundError:
    pass

Fix 7: Async context managers

import asyncio
from contextlib import asynccontextmanager

@asynccontextmanager
async def async_database():
    conn = await async_create_connection()
    try:
        yield conn
    finally:
        await conn.close()

# Usage
async def main():
    async with async_database() as conn:
        await conn.execute("SELECT * FROM users")

Lessons Learned

  • Use @contextmanager decorator for simple context managers — it is much cleaner than writing __enter__ and __exit__ methods.
  • Always put cleanup code in finally within the generator to ensure it runs even on exceptions.
  • Return False from __exit__ unless you specifically want to suppress exceptions.
  • Use contextlib.ExitStack when you need to manage a variable number of resources.
  • Build a library of reusable context managers for common patterns like database connections, file handling, and timing.

This blog does not accept any external sponsorships, affiliate marketing, or ad revenue.