Context Managers & Protocols

# Create the file so the three variants below have something to open
with open('data.txt', 'w') as f:
    f.write("hello")

# Bad: Resource leak if exception occurs
leaked = open('data.txt')
try:
    leaked.read()
    raise ValueError("something went wrong")
    leaked.close()  # This never executes!
except ValueError:
    pass
print(f"bad:    still open? {not leaked.closed}")

# Better but verbose
guarded = open('data.txt')
try:
    guarded.read()
    raise ValueError("something went wrong")
except ValueError:
    pass
finally:
    guarded.close()  # Guaranteed to run
print(f"better: still open? {not guarded.closed}")

# Best: Context manager handles everything
try:
    with open('data.txt') as managed:
        managed.read()
        raise ValueError("something went wrong")
except ValueError:
    pass
print(f"best:   still open? {not managed.closed}")

# Create the file first, so the example below has something to open
with open('data.txt', 'w') as f:
    f.write("hello from data.txt")


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()
    print(content)
    # File automatically closes, even if exception occurs

import os

from contextlib import suppress

# Standard library: suppress specific exceptions
with suppress(FileNotFoundError, PermissionError):
    os.remove('might_not_exist.txt')
print("still running after the missing file")

# A context manager cannot retry its own block. Returning True from
# __exit__ suppresses the exception and resumes AFTER the with statement,
# so the body still runs exactly once. Retrying needs a loop around it.
class SuppressFailure:
    """Swallow whatever the block raises, and record it"""

    def __init__(self):
        self.failed = False
        self.error = None

    def __enter__(self):
        self.failed = False
        self.error = None
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            return False  # Nothing to suppress
        self.failed = True
        self.error = exc_value
        return True  # Suppressed: control resumes after the with block

# Fails the first two times, then succeeds, so the retry loop is visible
tries = []

def risky_operation():
    tries.append(1)
    if len(tries) < 3:
        raise ConnectionError(f"connection refused (try {len(tries)})")
    print("risky_operation succeeded")

# The loop is what retries; the context manager just keeps it readable
attempt_state = SuppressFailure()
for attempt in range(1, 4):
    with attempt_state:
        risky_operation()
    if not attempt_state.failed:
        break
    print(f"Attempt {attempt} failed: {attempt_state.error}")
else:
    print("All 3 attempts failed")

# 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)

with IgnoreSpecificErrors(KeyError):
    {}['missing']
print("KeyError was ignored")

try:
    with IgnoreSpecificErrors(KeyError):
        raise TypeError("not in the ignore list")
except TypeError as e:
    print(f"TypeError propagated: {e}")

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(f"inside: debug={config.debug}")
print(f"restored: debug={config.debug}")

@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

# Usage
with atomic_write('config.json') as f:
    f.write('{"ok": true}')

with open('config.json') as f:
    print(f"config.json holds: {f.read()}")

# A failure leaves no half-written file behind
try:
    with atomic_write('report.json') as f:
        f.write('{"partial": ')
        raise ValueError("ran out of data")
except ValueError:
    pass
print(f"report.json exists: {os.path.exists('report.json')}")
print(f"report.json.tmp exists: {os.path.exists('report.json.tmp')}")

import json
import shutil
import tempfile
import threading
from contextlib import ExitStack, contextmanager
from pathlib import Path

work = Path(tempfile.mkdtemp())
(work / "a.json").write_text('{"debug": true}')
(work / "b.json").write_text('{"retries": 3}')

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 the loop fails partway
        return [len(f.read()) for f in files]

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()}

print(f"bytes read: {process_files([work / 'a.json', work / 'b.json'])}")
print(f"merged: {merge_configs(work / 'a.json', work / 'b.json')}")

# 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

lock = threading.Lock()
with conditional_lock(lock, should_lock=True):
    print(f"locked inside: {lock.locked()}")
print(f"locked after: {lock.locked()}")

with conditional_lock(lock, should_lock=False):
    print(f"never locked: {lock.locked()}")

# Callback registration for custom cleanup
with ExitStack() as stack:
    # Register arbitrary cleanup functions, run in reverse order on exit
    stack.callback(print, "Cleaning up...")
    stack.callback(shutil.rmtree, work, ignore_errors=True)

    # Do work...
    print(f"work dir exists: {work.exists()}")

print(f"work dir exists after: {work.exists()}")

from contextlib import contextmanager


class Connection:
    """Stand-in for a real DB connection so the example runs on its own"""

    def begin(self):
        print("BEGIN")

    def execute(self, sql):
        print(f"  {sql}")

    def commit(self):
        print("COMMIT")

    def rollback(self):
        print("ROLLBACK")


@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

db = Connection()

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

# The same manager rolls back when the block raises
try:
    with transaction(db) as conn:
        conn.execute("INSERT INTO orders ...")
        raise RuntimeError("constraint violated")
except RuntimeError as e:
    print(f"propagated: {e}")

import threading
from contextlib import contextmanager


class SharedResource:
    """Stand-in for whatever the critical section protects"""

    def __init__(self):
        self.value = 0

    def modify(self):
        self.value += 1


shared_resource = SharedResource()

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

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

print(f"value: {shared_resource.value}, still locked: {lock.locked()}")

# 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()

with timeout_lock(lock, 1.0):
    print(f"inside timeout_lock, locked: {lock.locked()}")

# Hold the lock elsewhere and the timeout fires instead of blocking forever
lock.acquire()
try:
    with timeout_lock(lock, 0.1):
        pass
except TimeoutError as e:
    print(f"TimeoutError: {e}")
finally:
    lock.release()

import os
import tempfile
from contextlib import contextmanager

@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

started_in = os.getcwd()
elsewhere = tempfile.mkdtemp()

with chdir(elsewhere):
    print(f"inside chdir:  {os.getcwd() == os.path.realpath(elsewhere)}")
print(f"back where we started: {os.getcwd() == started_in}")

print(f"DEPLOY_ENV set before: {'DEPLOY_ENV' in os.environ}")
with env_var('DEPLOY_ENV', 'staging'):
    print(f"inside env_var: {os.environ['DEPLOY_ENV']}")
print(f"DEPLOY_ENV set after:  {'DEPLOY_ENV' in os.environ}")

import logging
import time
from contextlib import contextmanager

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
logger = logging.getLogger("pipeline")

@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

def extract_data():
    pass

def transform_data():
    pass

def load_data():
    raise ConnectionError("warehouse unreachable")

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

# The failure path logs at ERROR and still re-raises
try:
    with log_execution("nightly load", logger):
        load_data()
except ConnectionError:
    # Everything here goes through logging, so there is one stream and one
    # ordering. Mixing print (stdout) with logging (stderr) would interleave
    # differently depending on whether you are on a terminal or a pipe.
    logger.warning("caller still sees the exception")

import io
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:
        return "used it"

# 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:
        return obj.read()
    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."""
    return stream.read().upper()

# BytesIO never inherits from any of these, it just has the methods
stream = io.BytesIO(b"payload")
print(f"BytesIO is Closeable: {isinstance(stream, Closeable)}")
print(f"int is Closeable: {isinstance(42, Closeable)}")
print(f"BytesIO is a context manager: {isinstance(stream, SupportsContextManager)}")

print(f"process_stream: {process_stream(stream)}")
print(f"use_resource: {use_resource(io.BytesIO(b'x'))}")

closed_it = io.BytesIO(b"drain me")
print(f"ensure_closed returned: {ensure_closed(closed_it)}")
print(f"and closed it: {closed_it.closed}")

import asyncio
import time
from contextlib import asynccontextmanager

# Class-based async context manager
class AsyncConnection:
    async def connect(self):
        await asyncio.sleep(0)

    async def close(self):
        await asyncio.sleep(0)

    async def query(self, sql):
        await asyncio.sleep(0)
        return f"2 rows from {sql}"

    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

# Generator-based async context manager
@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 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)

class Pool:
    """Stand-in for a real pool so the example runs on its own"""

    def __init__(self):
        self.in_use = 0

    async def acquire(self):
        self.in_use += 1
        return AsyncConnection()

    async def release(self, conn):
        self.in_use -= 1

class AsyncDB:
    async def begin(self):
        print("BEGIN")

    async def commit(self):
        print("COMMIT")

    async def rollback(self):
        print("ROLLBACK")

async def main():
    async with AsyncConnection() as conn:
        print(await conn.query("users"))

    async with async_timer("API call"):
        await asyncio.sleep(0.1)

    try:
        async with async_transaction(AsyncDB()):
            raise RuntimeError("deadlock detected")
    except RuntimeError as e:
        print(f"propagated: {e}")

    pool = Pool()
    async with get_connection(pool) as conn:
        print(f"connections in use: {pool.in_use}")
    print(f"connections in use after: {pool.in_use}")

asyncio.run(main())

import tempfile
from contextlib import contextmanager

# 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

path = tempfile.mkstemp()[1]

for factory in (bad_file_handler, good_file_handler):
    try:
        with factory(path) as f:
            raise ValueError("boom")
    except ValueError:
        pass
    print(f"{factory.__name__:18} left the file open: {not f.closed}")

# BAD: Suppressing all exceptions
class DangerousContext:
    def __enter__(self):
        return self

    def __exit__(self, *args):
        return True  # Silences KeyboardInterrupt, SystemExit!

# GOOD: Selective suppression
class SafeContext:
    def __enter__(self):
        return self

    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))

# A blanket "return True" swallows the one exception you must never swallow
with DangerousContext():
    raise KeyboardInterrupt("user pressed ctrl-c")
print("DangerousContext swallowed a KeyboardInterrupt")

with SafeContext():
    raise ValueError("expected, and on the list")
print("SafeContext suppressed the ValueError")

try:
    with SafeContext():
        raise KeyboardInterrupt("user pressed ctrl-c")
except KeyboardInterrupt:
    print("SafeContext let the KeyboardInterrupt through")

from contextlib import contextmanager

# BAD: Stateful generator context manager
@contextmanager
def bad_counter():
    count = 0
    def increment():
        nonlocal count
        count += 1
    yield increment
    print(f"Final count: {count}")

# 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

# One generator drives one 'with' block, so the object cannot be reused
cm = bad_counter()
with cm as inc:
    inc()
try:
    with cm as inc:
        inc()
except AttributeError as e:
    print(f"reuse failed: {type(e).__name__}: {e}")

# The class-based one carries its state and can be entered again
counter = Counter()
with counter:
    counter.increment()
with counter:
    counter.increment()