Context Managers and Resource Management
45 minContext managers are Python objects that define what happens when you enter and exit a block of code using the `with` statement. They ensure proper resource management by automatically handling setup and cleanup, even when exceptions occur. The `with` statement guarantees that cleanup code (in `__exit__`) runs regardless of how the block exits (normally or via exception). This makes context managers essential for managing resources like files, database connections, locks, network connections, and any resource that requires proper initialization and cleanup.
Context managers implement the context manager protocol with two methods: `__enter__()` (called when entering the `with` block, returns the resource) and `__exit__()` (called when exiting, handles cleanup and optionally suppresses exceptions). The `with` statement calls `__enter__()` at the start, assigns its return value to the variable after `as`, executes the block, and calls `__exit__()` when done (even if an exception occurs). This pattern ensures resources are always properly cleaned up, preventing leaks and errors.
Python's `contextlib` module provides utilities for creating context managers more easily. The `@contextmanager` decorator converts a generator function into a context manager—everything before `yield` runs in `__enter__()`, and everything after runs in `__exit__()`. `ExitStack` allows managing multiple context managers dynamically. `closing()` wraps objects with `close()` methods. `suppress()` suppresses specific exceptions. These utilities make it easier to create context managers without writing full classes.
Common use cases for context managers include file operations (ensuring files are closed), database connections (ensuring connections are closed and transactions are committed/rolled back), locks (ensuring locks are released), temporary resources (ensuring cleanup), and custom resource management (timers, performance monitoring, error handling). The `with` statement is so common in Python that many built-in types (like `open()` for files) are context managers, and libraries provide context managers for their resources.
Best practices include always using `with` statements for resource management, creating context managers for custom resources that need cleanup, using `contextlib` utilities when appropriate, handling exceptions properly in `__exit__` (return True to suppress, False to propagate), and documenting context manager behavior. Context managers make code more robust, readable, and maintainable by ensuring resources are always properly managed, even in error cases.
Advanced patterns include nested context managers (multiple `with` statements), context manager composition (managing multiple resources together), context manager factories (functions that return context managers), and using `ExitStack` for dynamic resource management. Understanding context managers is essential for writing professional Python code that handles resources correctly and avoids leaks and errors.
Key Concepts
- Context managers ensure proper resource setup and cleanup.
- The with statement automatically calls __enter__() and __exit__().
- Context managers guarantee cleanup even when exceptions occur.
- contextlib provides utilities for creating context managers easily.
- Context managers are essential for file, database, and lock management.
Learning Objectives
Master
- Creating context managers for resource management
- Using the with statement for automatic resource cleanup
- Working with contextlib utilities for easier context manager creation
- Implementing proper error handling in context managers
Develop
- Resource management thinking and best practices
- Understanding automatic cleanup and exception safety
- Designing robust, maintainable resource management code
Tips
- Always use with statements for file operations and resource management.
- Use @contextmanager decorator for simple context managers.
- Handle exceptions properly in __exit__ (return True to suppress).
- Use ExitStack for managing multiple context managers dynamically.
Common Pitfalls
- Not using with statements, causing resource leaks.
- Forgetting to handle exceptions in __exit__, causing unexpected behavior.
- Not understanding that __exit__ always runs, even on exceptions.
- Creating context managers that don't properly clean up resources.
Summary
- Context managers ensure proper resource setup and cleanup.
- The with statement automatically manages context manager lifecycle.
- Context managers guarantee cleanup even when exceptions occur.
- contextlib provides utilities for easier context manager creation.
- Using context managers makes code more robust and maintainable.
Exercise
Create a comprehensive context manager system for managing database connections, file operations, and custom resources with proper error handling and cleanup.
from contextlib import contextmanager, ExitStack
from typing import Optional, List, Dict, Any
import time
import logging
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class DatabaseConnection:
"""Simulated database connection class."""
def __init__(self, connection_string: str):
self.connection_string = connection_string
self.is_connected = False
self.transaction_count = 0
def connect(self):
"""Simulate connecting to database."""
logger.info(f"Connecting to database: {self.connection_string}")
time.sleep(0.1) # Simulate connection time
self.is_connected = True
logger.info("Database connected successfully")
def disconnect(self):
"""Simulate disconnecting from database."""
if self.is_connected:
logger.info("Disconnecting from database")
self.is_connected = False
logger.info("Database disconnected")
def execute_query(self, query: str) -> List[Dict[str, Any]]:
"""Simulate executing a database query."""
if not self.is_connected:
raise RuntimeError("Database not connected")
logger.info(f"Executing query: {query}")
time.sleep(0.05) # Simulate query execution time
# Return mock data
return [{"id": 1, "name": "Sample Data", "query": query}]
def begin_transaction(self):
"""Simulate beginning a transaction."""
if not self.is_connected:
raise RuntimeError("Database not connected")
self.transaction_count += 1
logger.info(f"Transaction {self.transaction_count} started")
def commit_transaction(self):
"""Simulate committing a transaction."""
if self.transaction_count > 0:
logger.info(f"Transaction {self.transaction_count} committed")
self.transaction_count -= 1
def rollback_transaction(self):
"""Simulate rolling back a transaction."""
if self.transaction_count > 0:
logger.info(f"Transaction {self.transaction_count} rolled back")
self.transaction_count -= 1
class FileManager:
"""Context manager for file operations with backup and logging."""
def __init__(self, filename: str, mode: str = 'r', backup: bool = True):
self.filename = filename
self.mode = mode
self.backup = backup
self.file = None
self.backup_filename = None
def __enter__(self):
"""Enter the context and prepare the file."""
# Create backup if needed
if self.backup and self.mode in ['w', 'a'] and self._file_exists():
self.backup_filename = f"{self.filename}.backup"
self._create_backup()
# Open the file
self.file = open(self.filename, self.mode)
logger.info(f"File {self.filename} opened in {self.mode} mode")
return self.file
def __exit__(self, exc_type, exc_val, exc_tb):
"""Exit the context and clean up resources."""
if self.file:
self.file.close()
logger.info(f"File {self.filename} closed")
# Handle exceptions
if exc_type is not None:
logger.error(f"Error occurred while working with {self.filename}: {exc_val}")
if self.backup_filename and self._file_exists(self.backup_filename):
self._restore_backup()
return False # Re-raise the exception
return True # Suppress the exception
def _file_exists(self, filename: str = None) -> bool:
"""Check if a file exists."""
import os
return os.path.exists(filename or self.filename)
def _create_backup(self):
"""Create a backup of the original file."""
import shutil
shutil.copy2(self.filename, self.backup_filename)
logger.info(f"Backup created: {self.backup_filename}")
def _restore_backup(self):
"""Restore file from backup."""
import shutil
shutil.copy2(self.backup_filename, self.filename)
logger.info(f"File restored from backup: {self.backup_filename}")
# Clean up backup
import os
os.remove(self.backup_filename)
logger.info(f"Backup file removed: {self.backup_filename}")
class PerformanceMonitor:
"""Context manager for monitoring performance of code blocks."""
def __init__(self, operation_name: str, log_memory: bool = True):
self.operation_name = operation_name
self.log_memory = log_memory
self.start_time = None
self.start_memory = None
def __enter__(self):
"""Enter the context and start monitoring."""
self.start_time = time.time()
if self.log_memory:
self.start_memory = self._get_memory_usage()
logger.info(f"Starting performance monitoring for: {self.operation_name}")
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Exit the context and log performance metrics."""
end_time = time.time()
execution_time = end_time - self.start_time
if self.log_memory and self.start_memory is not None:
end_memory = self._get_memory_usage()
memory_delta = end_memory - self.start_memory
logger.info(f"{self.operation_name} completed in {execution_time:.4f}s, "
f"memory delta: {memory_delta:+d} bytes")
else:
logger.info(f"{self.operation_name} completed in {execution_time:.4f}s")
return False # Don't suppress exceptions
def _get_memory_usage(self) -> int:
"""Get current memory usage."""
try:
import psutil
return psutil.Process().memory_info().rss
except ImportError:
return 0
class TransactionManager:
"""Context manager for database transactions."""
def __init__(self, connection: DatabaseConnection):
self.connection = connection
self.transaction_started = False
def __enter__(self):
"""Enter the context and start a transaction."""
self.connection.begin_transaction()
self.transaction_started = True
return self.connection
def __exit__(self, exc_type, exc_val, exc_tb):
"""Exit the context and handle transaction completion."""
if self.transaction_started:
if exc_type is None:
# No exception occurred, commit the transaction
self.connection.commit_transaction()
logger.info("Transaction committed successfully")
else:
# Exception occurred, rollback the transaction
self.connection.rollback_transaction()
logger.warning("Transaction rolled back due to exception")
return False # Re-raise the exception
return True
@contextmanager
def temporary_file(filename: str, content: str = ""):
"""Context manager for creating temporary files."""
temp_file = None
try:
temp_file = open(filename, 'w')
if content:
temp_file.write(content)
temp_file.flush()
yield filename
finally:
if temp_file:
temp_file.close()
# Clean up the temporary file
import os
if os.path.exists(filename):
os.remove(filename)
logger.info(f"Temporary file {filename} removed")
@contextmanager
def error_handler(operation_name: str, default_value: Any = None):
"""Context manager for handling errors gracefully."""
try:
yield
except Exception as e:
logger.error(f"Error in {operation_name}: {e}")
if default_value is not None:
logger.info(f"Returning default value: {default_value}")
return default_value
raise
# Example usage functions
def demonstrate_context_managers():
"""Demonstrate various context managers in action."""
print("=== Context Manager Demonstration ===\n")
# Database connection example
print("1. Database Connection Management:")
db = DatabaseConnection("postgresql://localhost:5432/mydb")
try:
with TransactionManager(db) as conn:
conn.connect()
result = conn.execute_query("SELECT * FROM users")
print(f" Query result: {result}")
# Simulate some work
time.sleep(0.1)
except Exception as e:
print(f" Database error: {e}")
finally:
db.disconnect()
# File management example
print("\n2. File Management with Backup:")
test_filename = "test_file.txt"
try:
with FileManager(test_filename, 'w', backup=True) as f:
f.write("This is test content\n")
f.write("Line 2\n")
print(f" File {test_filename} written successfully")
# Read the file back
with FileManager(test_filename, 'r') as f:
content = f.read()
print(f" File content: {repr(content)}")
except Exception as e:
print(f" File error: {e}")
# Performance monitoring example
print("\n3. Performance Monitoring:")
with PerformanceMonitor("Data Processing Operation"):
# Simulate some work
time.sleep(0.2)
data = [i ** 2 for i in range(10000)]
result = sum(data)
print(f" Processing result: {result}")
# Multiple context managers with ExitStack
print("\n4. Multiple Context Managers:")
with ExitStack() as stack:
# Open multiple files
files = []
for i in range(3):
filename = f"temp_file_{i}.txt"
file_obj = stack.enter_context(open(filename, 'w'))
file_obj.write(f"Content for file {i}\n")
files.append(filename)
print(f" Created {len(files)} temporary files")
# All files will be automatically closed when exiting
# Custom context manager with error handling
print("\n5. Error Handling Context Manager:")
with error_handler("Division Operation", default_value=0):
result = 10 / 0 # This will cause an error
print(f" Division result: {result}")
# This should work normally
with error_handler("Safe Operation"):
result = 10 / 2
print(f" Safe division result: {result}")
def demonstrate_resource_management():
"""Demonstrate advanced resource management patterns."""
print("\n=== Advanced Resource Management ===\n")
# Nested context managers
print("1. Nested Context Managers:")
with PerformanceMonitor("Database Operations"):
with TransactionManager(DatabaseConnection("sqlite:///test.db")) as db:
db.connect()
# Multiple operations within the transaction
for i in range(3):
result = db.execute_query(f"SELECT * FROM table_{i}")
print(f" Query {i} result: {len(result)} rows")
# Simulate some processing time
time.sleep(0.1)
# Context manager composition
print("\n2. Context Manager Composition:")
def combined_context_managers():
"""Combine multiple context managers."""
with PerformanceMonitor("Combined Operation"):
with error_handler("File Processing", default_value=""):
with FileManager("combined_test.txt", 'w') as f:
f.write("Combined context manager test\n")
return "Operation completed successfully"
result = combined_context_managers()
print(f" Combined operation result: {result}")
if __name__ == "__main__":
demonstrate_context_managers()
demonstrate_resource_management()