	In Browser
	StumbleUpon
	del.icio.us
	Google
	Google Buzz
	reddit
	LinkedIn

	Facebook
	Twitter
	Linkedin
	E-Mail

Modern Python > Type Hints and Dataclasses > Python Context Managers and the with Statement

Python Context Managers and the with Statement

Author: Venkata Sudhakar

A Python context manager is an object that defines a runtime context for use with the with statement. The with statement guarantees that setup code runs before the block and teardown code runs after - even if an exception occurs inside the block. The most familiar example is with open("file.txt") as f: which guarantees the file is closed when the block exits, regardless of whether an exception was raised. Context managers eliminate the try/finally boilerplate that would otherwise be needed for resource cleanup.

Any class that implements __enter__ and __exit__ methods is a context manager. __enter__ runs on entry to the with block and returns the value bound by the as clause. __exit__ runs on exit and receives exception information if one was raised - returning True suppresses the exception, returning False (or None) re-raises it. The contextlib module provides two easier ways to write context managers: the @contextmanager decorator for generator-based context managers, and contextlib.suppress() for ignoring specific exceptions.

The below example shows practical context managers for database connections, migration timing, and temporary directory management used in data migration pipelines.

import time
import contextlib
from typing import Optional

# Context manager 1: Database connection with automatic commit/rollback
class DatabaseTransaction:
    def __init__(self, connection_string: str):
        self.connection_string = connection_string
        self.conn = None
        self.cursor = None

def __enter__(self):
        print(f"  Connecting to {self.connection_string}...")
        # In production: self.conn = psycopg2.connect(self.connection_string)
        self.conn = type("MockConn", (), {
            "commit": lambda self: print("  COMMIT"),
            "rollback": lambda self: print("  ROLLBACK"),
            "cursor": lambda self: type("MockCursor", (), {
                "execute": lambda self, sql: print(f"  SQL: {sql[:50]}..."),
                "close": lambda self: None
            })()
        })()
        self.cursor = self.conn.cursor()
        return self.cursor  # bound to the "as" variable

def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.conn.commit()  # success: commit
            print("  Transaction committed.")
        else:
            self.conn.rollback()  # error: rollback
            print(f"  Transaction rolled back due to: {exc_val}")
        if self.cursor:
            self.cursor.close()
        return False  # do not suppress exceptions

# Context manager 2: Migration phase timer using @contextmanager
@contextlib.contextmanager
def migration_phase(phase_name: str):
    print(f"\n[START] {phase_name}")
    start = time.perf_counter()
    try:
        yield  # control passes to the with block here
    except Exception as e:
        elapsed = time.perf_counter() - start
        print(f"[FAILED] {phase_name} after {elapsed:.2f}s: {e}")
        raise
    else:
        elapsed = time.perf_counter() - start
        print(f"[DONE] {phase_name} in {elapsed:.2f}s")

# Using the context managers
print("=== Migration Transaction Example ===")
try:
    with DatabaseTransaction("postgresql://prod-db:5432/appdb") as cursor:
        cursor.execute("INSERT INTO migration_log VALUES (1, CURRENT_TIMESTAMP)")
        cursor.execute("UPDATE customers SET migrated=true WHERE id < 1000")
except Exception as e:
    print(f"Error: {e}")

print("\n=== Migration Phase Timing ===")
with migration_phase("Initial snapshot - customers table"):
    time.sleep(0.1)  # simulate work

with migration_phase("CDC sync validation"):
    time.sleep(0.05)  # simulate work

It gives the following output,

=== Migration Transaction Example ===
  Connecting to postgresql://prod-db:5432/appdb...
  SQL: INSERT INTO migration_log VALUES (1, CURRENT_TIMESTAMP)...
  SQL: UPDATE customers SET migrated=true WHERE id < 1000...
  COMMIT
  Transaction committed.

=== Migration Phase Timing ===

[START] Initial snapshot - customers table
[DONE] Initial snapshot - customers table in 0.10s

[START] CDC sync validation
[DONE] CDC sync validation in 0.05s

import tempfile
import os
import contextlib

# Context manager 3: Temporary working directory for migration staging
@contextlib.contextmanager
def temp_migration_workspace(prefix: str = "migration_"):
    tmp_dir = tempfile.mkdtemp(prefix=prefix)
    print(f"Created workspace: {tmp_dir}")
    try:
        yield tmp_dir  # caller gets the temp directory path
    finally:
        # Always clean up, even on exception
        import shutil
        shutil.rmtree(tmp_dir, ignore_errors=True)
        print(f"Cleaned up workspace: {tmp_dir}")

# Context manager 4: contextlib.suppress - ignore specific non-critical errors
def check_table_exists(table_name: str) -> bool:
    with contextlib.suppress(KeyError):
        tables = {"customers": True, "orders": True}
        return tables[table_name]  # raises KeyError if not found
    return False  # KeyError suppressed, return default

# Nested context managers on one line (Python 3.10+ supports cleaner syntax)
with (temp_migration_workspace("cdc_export_") as workspace,
      migration_phase("Export CDC events to staging")):
    # Write staging files to the temp directory
    staging_file = os.path.join(workspace, "cdc_events.json")
    with open(staging_file, "w") as f:
        f.write("{events: 125000}")
    print(f"Staged {os.path.getsize(staging_file)} bytes")

print()
print("customers exists:", check_table_exists("customers"))
print("inventory exists:", check_table_exists("inventory"))  # suppressed KeyError

It gives the following output,

Created workspace: /tmp/cdc_export_abc123

[START] Export CDC events to staging
Staged 17 bytes
[DONE] Export CDC events to staging in 0.00s
Cleaned up workspace: /tmp/cdc_export_abc123

customers exists: True
inventory exists: False

When to use context managers:

Use context managers whenever you have setup and teardown logic that must be paired: opening/closing files, acquiring/releasing locks, starting/committing database transactions, beginning/ending timers, creating/deleting temporary resources, and enabling/disabling features for a block of code. The with statement makes the pairing explicit and exception-safe. Prefer @contextlib.contextmanager for simple cases - it avoids the class boilerplate and makes the setup/teardown flow visually clear because the yield clearly separates the two phases.

Send your comments, suggestions or queries regarding this site to [email protected].