Python Context Managers and with Statement: Clean Resource Management
Master Python context managers for safe resource handling. Learn with statement, __enter__/__exit__, contextlib, and best practices for file handling and database connections
📅 Published: May 25, 2025 ✏️ Updated: July 5, 2025 By Ojaswi Athghara
#python #context #with #resources #cleanup
The Bug That Taught Me Context Managers
My production ML script kept crashing with "too many open files." I'd forgotten to close files:
# My buggy code
file = open('data.csv')
data = file.read()
# Forgot file.close()!
After 1024 files, the system crashed. Context managers saved me:
# Automatic cleanup!
with open('data.csv') as file:
data = file.read()
# File automatically closed!
Context managers = Automatic resource cleanup!
Understanding with Statement
The with statement ensures resources are properly cleaned up, even if errors occur.
Without Context Manager (Dangerous)
# Manual resource management - risky!
file = open('data.txt', 'r')
try:
data = file.read()
# Process data...
finally:
file.close() # Must remember to close!
With Context Manager (Safe)
# Automatic resource management - safe!
with open('data.txt', 'r') as file:
data = file.read()
# Process data...
# File automatically closed, even if error occurs!
Creating Custom Context Managers
Method 1: Class-Based
Class-based context managers implement __enter__ and __exit__ methods:
class DatabaseConnection:
def __init__(self, db_name):
self.db_name = db_name
self.connection = None
def __enter__(self):
"""Called when entering with block."""
print(f"Opening connection to {self.db_name}")
# Establish connection
self.connection = f"Connection to {self.db_name}"
# Return value is assigned to 'as' variable
return self.connection
def __exit__(self, exc_type, exc_val, exc_tb):
"""Called when exiting with block."""
print(f"Closing connection to {self.db_name}")
self.connection = None
# exc_type: Exception class (if exception occurred)
# exc_val: Exception instance
# exc_tb: Traceback object
# Return False to propagate exceptions
# Return True to suppress exceptions
return False
# Usage
with DatabaseConnection('mydb') as conn:
print(f"Using {conn}")
# Do database operations
# Connection automatically closed
Understanding exit Parameters
The __exit__ method receives three arguments when an exception occurs:
class ErrorHandler:
def __enter__(self):
print("Entering context")
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
print("No errors occurred")
else:
print(f"Exception type: {exc_type}")
print(f"Exception value: {exc_val}")
print(f"Traceback: {exc_tb}")
# Return True to suppress the exception
# Return False (or None) to propagate it
return False
# Example with error
try:
with ErrorHandler():
print("Inside context")
raise ValueError("Something went wrong!")
except ValueError as e:
print(f"Caught: {e}")
Exception Suppression
class ErrorSuppressor:
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is ValueError:
print(f"Suppressing ValueError: {exc_val}")
return True # Suppress this specific exception
return False # Propagate other exceptions
# ValueError is suppressed
with ErrorSuppressor():
raise ValueError("This won't crash the program")
print("Program continues...")
# Other exceptions propagate
try:
with ErrorSuppressor():
raise TypeError("This will crash")
except TypeError:
print("TypeError was not suppressed")
Method 2: Using contextlib
The @contextmanager decorator makes creating context managers simpler:
from contextlib import contextmanager
@contextmanager
def timer(name):
"""Time a code block."""
import time
start = time.time()
print(f"Starting {name}...")
try:
yield # Code block runs here
finally:
# Cleanup code - always runs
end = time.time()
print(f"{name} took {end - start:.2f} seconds")
# Usage
with timer("Data processing"):
# Expensive operation
sum(range(1000000))
How it works:
- Code before
yieldruns in__enter__ yieldreturns control to thewithblock- Code after
yieldruns in__exit__
Yielding Values
@contextmanager
def temp_file():
"""Create and cleanup a temporary file."""
import tempfile
import os
# Create temp file
fd, path = tempfile.mkstemp()
print(f"Created temporary file: {path}")
try:
yield path # Provide path to with block
finally:
# Cleanup
os.close(fd)
os.unlink(path)
print(f"Deleted temporary file: {path}")
# Usage
with temp_file() as filepath:
# Write to the temp file
with open(filepath, 'w') as f:
f.write("Temporary data")
# Read it back
with open(filepath, 'r') as f:
print(f"Contents: {f.read()}")
# File is automatically deleted here
Handling Exceptions in contextlib
@contextmanager
def safe_division():
"""Handle division errors gracefully."""
try:
yield
except ZeroDivisionError:
print("Cannot divide by zero! Returning None")
# Exception is suppressed
except Exception as e:
print(f"Unexpected error: {e}")
raise # Re-raise other exceptions
# ZeroDivisionError is handled
with safe_division():
result = 10 / 0 # No crash!
print("Program continues")
Real-World Examples
File Operations
# Multiple files safely
with open('input.txt') as infile, open('output.txt', 'w') as outfile:
for line in infile:
outfile.write(line.upper())
Database Connections
@contextmanager
def get_db_connection(db_path):
import sqlite3
conn = sqlite3.connect(db_path)
try:
yield conn
conn.commit() # Commit if successful
except Exception:
conn.rollback() # Rollback on error
raise
finally:
conn.close() # Always close connection
# Usage
with get_db_connection('data.db') as conn:
cursor = conn.cursor()
cursor.execute('SELECT * FROM users')
results = cursor.fetchall()
Changing Directory Temporarily
@contextmanager
def change_dir(path):
"""Temporarily change working directory."""
import os
old_dir = os.getcwd()
os.chdir(path)
try:
yield
finally:
os.chdir(old_dir)
# Usage
print(f"Current dir: {os.getcwd()}")
with change_dir('/tmp'):
print(f"Temp dir: {os.getcwd()}")
# Do work in /tmp
print(f"Back to: {os.getcwd()}")
Thread Lock Management
import threading
@contextmanager
def acquire_lock(lock, timeout=10):
"""Safely acquire and release a thread lock."""
acquired = lock.acquire(timeout=timeout)
if not acquired:
raise RuntimeError(f"Could not acquire lock within {timeout}s")
try:
yield
finally:
lock.release()
# Usage
lock = threading.Lock()
with acquire_lock(lock):
# Critical section - thread-safe
shared_resource += 1
ML Model Training with Monitoring
@contextmanager
def training_monitor(model_name):
"""Monitor ML model training."""
import time
print(f"🚀 Starting training: {model_name}")
start = time.time()
try:
yield
except Exception as e:
print(f"❌ Training failed: {e}")
raise
else:
duration = time.time() - start
print(f"✅ Training complete in {duration:.1f}s")
# Usage
with training_monitor("RandomForest"):
# Train model
model.fit(X_train, y_train)
Setting Environment Variables Temporarily
@contextmanager
def temp_env_var(key, value):
"""Temporarily set an environment variable."""
import os
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
# Usage
with temp_env_var('DEBUG', 'true'):
# Environment has DEBUG=true
run_with_debug()
# DEBUG restored to original value
Advanced contextlib Utilities
suppress - Ignore Specific Exceptions
from contextlib import suppress
# Without suppress
try:
os.remove('file.txt')
except FileNotFoundError:
pass # Ignore if file doesn't exist
# With suppress (cleaner)
with suppress(FileNotFoundError):
os.remove('file.txt')
# Suppress multiple exception types
with suppress(ValueError, TypeError):
risky_operation()
redirect_stdout and redirect_stderr
from contextlib import redirect_stdout, redirect_stderr
import io
# Capture stdout
output = io.StringIO()
with redirect_stdout(output):
print("This goes to the StringIO")
print("Not to the console!")
print(f"Captured: {output.getvalue()}")
# Redirect to file
with open('output.log', 'w') as f:
with redirect_stdout(f):
print("This goes to output.log")
ExitStack - Dynamic Context Managers
from contextlib import ExitStack
# Manage unknown number of resources
def process_files(filenames):
with ExitStack() as stack:
# Open all files dynamically
files = [stack.enter_context(open(fname)) for fname in filenames]
# Process all files
for f in files:
process(f.read())
# All files automatically closed
# Another example: conditional resource management
with ExitStack() as stack:
if need_file:
f = stack.enter_context(open('data.txt'))
if need_db:
conn = stack.enter_context(get_db_connection('db.sqlite'))
# Use resources that were opened
closing - Ensure Objects Are Closed
from contextlib import closing
from urllib.request import urlopen
# Ensure urlopen is closed (doesn't support context manager natively in older Python)
with closing(urlopen('http://example.com')) as page:
content = page.read()
# page.close() automatically called
Common Pitfalls
Mistake 1: Forgetting to Return from enter
# Wrong - __enter__ returns None
class BadContext:
def __enter__(self):
self.resource = "Important data"
# Forgot to return self!
def __exit__(self, *args):
pass
with BadContext() as ctx:
print(ctx) # None - not what we wanted!
# Correct
class GoodContext:
def __enter__(self):
self.resource = "Important data"
return self # Return self or the resource
def __exit__(self, *args):
pass
with GoodContext() as ctx:
print(ctx.resource) # Works!
Mistake 2: Not Handling Exceptions in exit
# Dangerous - cleanup might fail
class RiskyContext:
def __exit__(self, *args):
self.file.close() # What if file doesn't exist?
# Safe - handle cleanup errors
class SafeContext:
def __exit__(self, *args):
try:
if hasattr(self, 'file'):
self.file.close()
except Exception as e:
print(f"Cleanup error: {e}")
Mistake 3: Modifying Exception Traceback
# Bad - loses exception information
class BadExceptionHandler:
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type:
# Creating new exception loses original traceback
raise ValueError("Something went wrong")
# Good - preserve traceback
class GoodExceptionHandler:
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type:
# Let original exception propagate
return False
Best Practices
1. Always Use with for File Operations
# Bad - easy to forget cleanup
file = open('data.txt')
data = file.read()
file.close() # Might not run if error occurs
# Good - automatic cleanup
with open('data.txt') as file:
data = file.read()
# Automatically closed, even if error occurs
2. Create Context Managers for Paired Operations
Any "setup-teardown" pattern is perfect for context managers:
# Setup-teardown pairs:
# - open → close
# - lock → unlock
# - start → stop
# - connect → disconnect
# - allocate → deallocate
@contextmanager
def performance_timer():
"""Start timer → measure elapsed → report."""
import time
start = time.time()
yield
print(f"Elapsed: {time.time() - start:.2f}s")
3. Use contextlib for Simple Cases
# Class-based (overkill for simple cases)
class SimpleTimer:
def __enter__(self):
self.start = time.time()
return self
def __exit__(self, *args):
print(f"Took {time.time() - self.start:.2f}s")
# contextlib (simpler!)
@contextmanager
def simple_timer():
start = time.time()
yield
print(f"Took {time.time() - start:.2f}s")
4. Document Your Context Managers
@contextmanager
def database_transaction(conn):
"""
Context manager for database transactions.
Automatically commits on success, rolls back on error.
Args:
conn: Database connection object
Yields:
Database cursor
Example:
with database_transaction(conn) as cursor:
cursor.execute("INSERT INTO users VALUES (?)", (name,))
"""
cursor = conn.cursor()
try:
yield cursor
conn.commit()
except Exception:
conn.rollback()
raise
finally:
cursor.close()
5. Handle Cleanup Exceptions Carefully
@contextmanager
def safe_resource():
"""Always prefer safe cleanup."""
resource = acquire_resource()
try:
yield resource
finally:
# Don't let cleanup errors mask original errors
try:
release_resource(resource)
except Exception as e:
# Log but don't raise during cleanup
logging.error(f"Cleanup failed: {e}")
Nesting Context Managers
Multiple Resources in One Statement
# Open multiple files at once
with open('input.txt') as infile, open('output.txt', 'w') as outfile:
for line in infile:
outfile.write(line.upper())
# Both files closed automatically
Nested with Statements
# Nested approach (more verbose)
with open('input.txt') as infile:
with open('output.txt', 'w') as outfile:
data = infile.read()
outfile.write(data.upper())
Combining Different Context Managers
with timer("Database operation"):
with get_db_connection('data.db') as conn:
with acquire_lock(db_lock):
# All three context managers active
cursor = conn.cursor()
cursor.execute("SELECT * FROM users")
# Released in reverse order: lock → connection → timer
Common Mistakes to Avoid
1. Forgetting to Return a Value
@contextmanager
def my_context():
setup()
# Missing yield!
cleanup()
# Error: generator didn't yield
2. Not Handling Exceptions in exit
# Bad - exceptions during cleanup go uncaught
def __exit__(self, exc_type, exc_val, exc_tb):
self.resource.close() # What if this fails?
return False
# Good - handle cleanup errors
def __exit__(self, exc_type, exc_val, exc_tb):
try:
self.resource.close()
except Exception as e:
logging.error(f"Cleanup error: {e}")
return False # Don't suppress original exception
3. Returning True Accidentally
Returning True from __exit__ suppresses exceptions, which can hide bugs! Only do this intentionally.
Key Takeaways
- Context managers guarantee cleanup - resources are released even if errors occur
- Use
withstatement for any resource that needs cleanup (files, locks, connections) - Two ways to create: class-based (
__enter__/__exit__) or@contextmanagerdecorator __exit__receives exception info - returnTrueto suppress,Falseto propagate- contextlib provides utilities:
suppress,redirect_stdout,ExitStack,closing - Perfect for paired operations: setup/teardown, start/stop, acquire/release
- Nest context managers for managing multiple resources safely
Context managers are Python's elegant solution to resource management. They make code cleaner, safer, and more maintainable. Start using them wherever you see setup-teardown patterns!
If this guide helped you understand context managers and the with statement, I'd love to hear about it! Connect with me on Twitter or LinkedIn.
Support My Work
If this guide helped you understand context managers, write cleaner Python code with the with statement, or build custom resource managers, I'd really appreciate your support! Creating comprehensive, beginner-friendly Python tutorials like this takes significant time and effort. Your support helps me continue sharing knowledge and creating more helpful resources for aspiring Python developers.
☕ Buy me a coffee - Every contribution, big or small, means the world to me and keeps me motivated to create more content!
Cover image by Nourhan Sabek on Unsplash