Context Managers & Protocols

# Bad: Resource leak if exception occurs
f = open('data.txt')
data = f.read()  # What if this raises an exception?
f.close()  # This never executes!

# Better but verbose
f = open('data.txt')
try:
    data = f.read()
finally:
    f.close()  # Guaranteed to run

# Best: Context manager handles everything
with open('data.txt') as f:
    data = f.read()  # Cleanup automatic, exception-safe

__enter__(self) -> Any
    # Called when entering 'with' block
    # Return value is bound to 'as' variable
    
__exit__(self, exc_type, exc_value, traceback) -> bool | None
    # Called when exiting 'with' block
    # Receives exception info if one occurred
    # Return True to suppress exception, False/None to propagate

class ManagedFile:
    def __init__(self, path, mode='r'):
        self.path = path
        self.mode = mode
        self.file = None

    def __enter__(self):
        print(f"Opening {self.path}")
        self.file = open(self.path, self.mode)
        return self.file  # This is what 'as f' receives

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"Closing {self.path}")
        if self.file:
            self.file.close()
        
        # Log exception if one occurred
        if exc_type is not None:
            print(f"Exception occurred: {exc_type.__name__}: {exc_value}")
        
        return False  # Propagate exceptions (don't suppress)

# Usage
with ManagedFile('data.txt') as f:
    content = f.read()
    # File automatically closes, even if exception occurs

from contextlib import suppress

# Standard library: suppress specific exceptions
with suppress(FileNotFoundError, PermissionError):
    os.remove('might_not_exist.txt')
# Continues execution if file doesn't exist

# Custom suppression for retryable operations
class RetryableOperation:
    def __init__(self, max_attempts=3):
        self.max_attempts = max_attempts
        self.attempt = 0
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            return False  # Success, no exception
        
        self.attempt += 1
        if self.attempt < self.max_attempts:
            print(f"Attempt {self.attempt} failed, retrying...")
            return True  # Suppress exception, allow retry
        
        print(f"All {self.max_attempts} attempts failed")
        return False  # Propagate exception

# Selective suppression by exception type
class IgnoreSpecificErrors:
    def __init__(self, *exception_types):
        self.exception_types = exception_types
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        # Suppress only specified exceptions
        return exc_type is not None and issubclass(exc_type, self.exception_types)

from contextlib import contextmanager
import time
import os

@contextmanager
def timer(name):
    """Time a code block and print duration"""
    start = time.time()
    print(f"Starting {name}...")
    try:
        yield  # Control returns to 'with' block here
    finally:
        duration = time.time() - start
        print(f"{name} took {duration:.2f} seconds")

# Usage
with timer("Data processing"):
    # Your code here
    time.sleep(1)

@contextmanager
def temporary_change(obj, attr, new_value):
    """Temporarily change an attribute, then restore it"""
    original = getattr(obj, attr)
    setattr(obj, attr, new_value)
    try:
        yield obj
    finally:
        setattr(obj, attr, original)

# Usage
from types import SimpleNamespace
config = SimpleNamespace(debug=False)
with temporary_change(config, 'debug', True):
    print(config.debug)  # True
print(config.debug)  # False (restored)

@contextmanager
def atomic_write(path):
    """Write to temp file, rename on success"""
    temp_path = f"{path}.tmp"
    f = open(temp_path, 'w')
    try:
        yield f
        f.close()
        os.rename(temp_path, path)  # Atomic on Unix
    except Exception:
        f.close()
        os.remove(temp_path)
        raise

import json
import shutil
from contextlib import ExitStack, contextmanager

def process_files(file_paths):
    """Process multiple files with guaranteed cleanup"""
    with ExitStack() as stack:
        # Dynamically add context managers
        files = [stack.enter_context(open(p)) for p in file_paths]

        # All files guaranteed to close, even if loop fails partway
        for f in files:
            data = f.read()
            # ... process data ...

def merge_configs(*config_paths):
    """Merge multiple config files safely"""
    with ExitStack() as stack:
        configs = [
            json.load(stack.enter_context(open(p)))
            for p in config_paths
        ]
        return {k: v for config in configs for k, v in config.items()}

# Conditional resource management
@contextmanager
def conditional_lock(lock, should_lock):
    """Optionally acquire a lock"""
    with ExitStack() as stack:
        if should_lock:
            stack.enter_context(lock)
        yield

# Callback registration for custom cleanup
temp_dir = "/tmp/work"
with ExitStack() as stack:
    # Register arbitrary cleanup functions
    stack.callback(print, "Cleaning up...")
    stack.callback(shutil.rmtree, temp_dir, ignore_errors=True)

    # Do work...

@contextmanager
def transaction(db_connection):
    """ACID transaction with automatic rollback on error"""
    db_connection.begin()
    try:
        yield db_connection
        db_connection.commit()
    except Exception:
        db_connection.rollback()
        raise

# Usage
with transaction(db) as conn:
    conn.execute("INSERT INTO users ...")
    conn.execute("UPDATE accounts ...")
    # Automatically commits if successful, rolls back if error

import threading

# Standard library locks are context managers
lock = threading.Lock()

with lock:
    # Critical section - guaranteed to release lock
    shared_resource.modify()

# Timeout-aware locking
@contextmanager
def timeout_lock(lock, timeout):
    acquired = lock.acquire(timeout=timeout)
    if not acquired:
        raise TimeoutError("Could not acquire lock")
    try:
        yield
    finally:
        lock.release()

@contextmanager
def chdir(path):
    """Temporarily change working directory"""
    original = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(original)

@contextmanager
def env_var(key, value):
    """Temporarily set environment variable"""
    old_value = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if old_value is None:
            del os.environ[key]
        else:
            os.environ[key] = old_value

@contextmanager
def log_execution(operation_name, logger):
    """Log operation execution with timing"""
    logger.info(f"Starting {operation_name}")
    start = time.time()
    try:
        yield
        duration = time.time() - start
        logger.info(f"{operation_name} succeeded in {duration:.2f}s")
    except Exception as e:
        duration = time.time() - start
        logger.error(f"{operation_name} failed after {duration:.2f}s: {e}")
        raise

with log_execution("ETL pipeline", logger):
    extract_data()
    transform_data()
    load_data()

from typing import Protocol, runtime_checkable

@runtime_checkable
class SupportsContextManager(Protocol):
    """Protocol for context manager objects"""
    def __enter__(self) -> object: ...
    def __exit__(self, exc_type, exc_value, traceback) -> bool | None: ...

# Any object with these methods is a context manager
def use_resource(resource: SupportsContextManager):
    with resource:
        # Use resource
        pass

# Custom protocol for closeable objects
@runtime_checkable
class Closeable(Protocol):
    def close(self) -> None: ...

def ensure_closed(obj: Closeable):
    """Works with files, sockets, database connections, etc."""
    try:
        # Use obj
        pass
    finally:
        obj.close()

# Protocol for readable objects
class Readable(Protocol):
    def read(self, size: int = -1) -> bytes: ...

def process_stream(stream: Readable):
    """Works with files, network streams, BytesIO, etc."""
    data = stream.read()
    return process(data)

# Class-based async context manager
class AsyncConnection:
    async def __aenter__(self):
        await self.connect()
        print("Connection established")
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        await self.close()
        print("Connection closed")
        return False

# Usage
async with AsyncConnection() as conn:
    await conn.query("SELECT * FROM users")

# Generator-based async context manager
from contextlib import asynccontextmanager

@asynccontextmanager
async def async_timer(name):
    """Async version of timer"""
    start = time.time()
    print(f"Starting {name}...")
    try:
        yield
    finally:
        duration = time.time() - start
        print(f"{name} took {duration:.2f} seconds")

async with async_timer("API call"):
    response = await fetch_data()

# Async database transaction
@asynccontextmanager
async def async_transaction(db):
    await db.begin()
    try:
        yield db
        await db.commit()
    except Exception:
        await db.rollback()
        raise

# Real-world example: connection pooling
@asynccontextmanager
async def get_connection(pool):
    """Get connection from pool, return on exit"""
    conn = await pool.acquire()
    try:
        yield conn
    finally:
        await pool.release(conn)

async with get_connection(db_pool) as conn:
    result = await conn.execute("SELECT ...")

# BAD: Not exception-safe
@contextmanager
def bad_file_handler(path):
    f = open(path)
    yield f
    f.close()  # Skipped if exception in with block!

# GOOD: Always use try/finally
@contextmanager
def good_file_handler(path):
    f = open(path)
    try:
        yield f
    finally:
        f.close()  # Always executes

# BAD: Suppressing all exceptions
class DangerousContext:
    def __exit__(self, *args):
        return True  # Silences KeyboardInterrupt, SystemExit!

# GOOD: Selective suppression
class SafeContext:
    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            return False
        # Only suppress expected exceptions
        return issubclass(exc_type, (ValueError, KeyError))

# BAD: Stateful generator context manager
@contextmanager
def bad_counter():
    count = 0
    def increment():
        nonlocal count
        count += 1
    yield increment
    print(f"Final count: {count}")
# Doesn't work - generator is consumed after first use!

# GOOD: Use class for stateful context managers
class Counter:
    def __init__(self):
        self.count = 0
    
    def __enter__(self):
        return self
    
    def increment(self):
        self.count += 1
    
    def __exit__(self, *args):
        print(f"Final count: {self.count}")
        return False