Context Managers (Custom)

# Class-based context manager
class DatabaseConnection:
    def __init__(self, db_name):
        self.db_name = db_name
        self.connected = False
    
    def __enter__(self):
        self.connected = True
        print(f"  Connected to {self.db_name}")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.connected = False
        print(f"  Disconnected from {self.db_name}")
        if exc_type:
            print(f"  Error: {exc_val}")
        return False  # Don't suppress exceptions

with DatabaseConnection("users_db") as db:
    print(f"  Is connected: {db.connected}")
    print("  Querying users...")

# contextlib decorator
from contextlib import contextmanager

@contextmanager
def temporary_directory(name):
    print(f"  Creating temp dir: {name}")
    try:
        yield name
    finally:
        print(f"  Cleaning up temp dir: {name}")

with temporary_directory("temp_data") as d:
    print(f"  Working in {d}")

# Practical: redirect output
import io
from contextlib import redirect_stdout

buffer = io.StringIO()
with redirect_stdout(buffer):
    print("This goes to buffer")
    print("Not to screen")

captured = buffer.getvalue()
print(f"Captured: {captured}")