DATATWEETS
Courses
Get Started

Lesson 9 - Context Managers and Resource Management

Introduction

You’ve probably seen code like this:

with open('books.csv', 'r') as file:
    content = file.read()

The with statement ensures the file is properly closed, even if an error occurs. This is an example of a context manager—an object that sets up and tears down resources automatically.

Context managers are essential for:

  • File operations
  • Database connections
  • Network connections
  • Locks and synchronization
  • Temporary state changes

In this lesson, we’ll learn:

  • How the with statement works
  • Creating context managers with classes (__enter__ and __exit__)
  • Creating context managers with @contextmanager decorator
  • Practical context manager examples
  • Nested context managers

The Problem: Manual Resource Management

Without context managers, you have to manually clean up resources:

# Risky code - what if an exception occurs?
file = open('books.csv', 'r')
content = file.read()
# Process content...
file.close()  # Might never execute if exception occurs!

You should use try/finally:

file = open('books.csv', 'r')
try:
    content = file.read()
    # Process content...
finally:
    file.close()  # Always executes

This works but is verbose and easy to forget.

The Solution: Context Managers

Context managers handle setup and cleanup automatically:

with open('books.csv', 'r') as file:
    content = file.read()
    # Process content...
# File is automatically closed here, even if an exception occurred

Much cleaner!

How Context Managers Work

The with statement calls two special methods:

  • __enter__(): Called when entering the with block
  • __exit__(): Called when leaving the with block (even if an exception occurs)

Here’s what happens:

with open('books.csv', 'r') as file:
    content = file.read()

# Equivalent to:
file = open('books.csv', 'r')
file_obj = file.__enter__()  # Returns self (the file object)
try:
    content = file_obj.read()
finally:
    file.__exit__(None, None, None)  # Always called

Creating a Simple Context Manager

Let’s create a context manager that prints messages on enter/exit:

class BookstoreConnection:
    """Simulate a database connection"""

    def __init__(self, connection_string):
        self.connection_string = connection_string
        self.connected = False

    def __enter__(self):
        """Setup: Connect to database"""
        print(f"Connecting to {self.connection_string}...")
        self.connected = True
        return self  # Return value is assigned to 'as' variable

    def __exit__(self, exc_type, exc_value, exc_traceback):
        """Cleanup: Disconnect from database"""
        print("Closing connection...")
        self.connected = False
        # Return False to propagate exceptions (default)
        return False

    def query(self, sql):
        """Execute a query"""
        if not self.connected:
            raise RuntimeError("Not connected!")
        print(f"Executing: {sql}")
        return ["Book 1", "Book 2", "Book 3"]

# Use the context manager
with BookstoreConnection("bookstore.db") as db:
    results = db.query("SELECT * FROM books")
    print(f"Results: {results}")

Output:

Connecting to bookstore.db...
Executing: SELECT * FROM books
Results: ['Book 1', 'Book 2', 'Book 3']
Closing connection...

Handling Exceptions in __exit__

The __exit__ method receives exception information if an error occurs:

class ErrorHandler:
    def __enter__(self):
        print("Entering context")
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        if exc_type is not None:
            print(f"An error occurred: {exc_type.__name__}: {exc_value}")
            # Return True to suppress the exception
            return True
        print("Exiting normally")
        return False

# Example 1: No error
print("=== No Error ===")
with ErrorHandler():
    print("Doing some work...")

# Example 2: With error (suppressed)
print("\n=== With Error (Suppressed) ===")
with ErrorHandler():
    print("Doing some work...")
    raise ValueError("Something went wrong!")
print("Program continues!")

Output:

=== No Error ===
Entering context
Doing some work...
Exiting normally

=== With Error (Suppressed) ===
Entering context
Doing some work...
An error occurred: ValueError: Something went wrong!
Program continues!

Practical Example: Transaction Manager

Let’s create a context manager for database-like transactions:

class Transaction:
    """Manage database transactions"""

    def __init__(self, db_name):
        self.db_name = db_name
        self.changes = []

    def __enter__(self):
        print(f"BEGIN TRANSACTION on {self.db_name}")
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        if exc_type is None:
            # No exception - commit
            print("COMMIT TRANSACTION")
            print(f"Applied {len(self.changes)} changes")
        else:
            # Exception occurred - rollback
            print("ROLLBACK TRANSACTION")
            print(f"Discarded {len(self.changes)} changes")
            self.changes.clear()

        # Don't suppress exceptions
        return False

    def add_book(self, title):
        """Add a book (simulated change)"""
        self.changes.append(('INSERT', title))
        print(f"  Added: {title}")

    def remove_book(self, title):
        """Remove a book (simulated change)"""
        self.changes.append(('DELETE', title))
        print(f"  Removed: {title}")

# Example 1: Successful transaction
print("=== Successful Transaction ===")
with Transaction("bookstore.db") as txn:
    txn.add_book("Python Basics")
    txn.add_book("Python Advanced")

# Example 2: Failed transaction (automatically rolls back)
print("\n=== Failed Transaction ===")
try:
    with Transaction("bookstore.db") as txn:
        txn.add_book("Data Science 101")
        txn.add_book("Machine Learning")
        raise RuntimeError("Network error!")
except RuntimeError:
    print("Handled the error")

Output:

=== Successful Transaction ===
BEGIN TRANSACTION on bookstore.db
  Added: Python Basics
  Added: Python Advanced
COMMIT TRANSACTION
Applied 2 changes

=== Failed Transaction ===
BEGIN TRANSACTION on bookstore.db
  Added: Data Science 101
  Added: Machine Learning
ROLLBACK TRANSACTION
Discarded 2 changes
Handled the error

The contextlib Module

Python’s contextlib module provides tools for working with context managers.

The @contextmanager Decorator

Instead of creating a class, you can use a generator function:

from contextlib import contextmanager

@contextmanager
def bookstore_connection(connection_string):
    """Context manager using a generator"""
    # Setup (before yield)
    print(f"Connecting to {connection_string}...")
    conn = {"connected": True, "name": connection_string}

    try:
        yield conn  # Value passed to 'as' variable
    finally:
        # Cleanup (after yield)
        print("Closing connection...")
        conn["connected"] = False

# Use it
with bookstore_connection("bookstore.db") as db:
    print(f"Using connection: {db}")

Output:

Connecting to bookstore.db...
Using connection: {'connected': True, 'name': 'bookstore.db'}
Closing connection...

This is much simpler than creating a class!

Practical Example: Timing Context Manager

from contextlib import contextmanager
import time

@contextmanager
def timer(name):
    """Time a block of code"""
    start_time = time.time()
    print(f"Starting {name}...")

    try:
        yield
    finally:
        end_time = time.time()
        duration = end_time - start_time
        print(f"{name} took {duration:.4f} seconds")

# Use it
with timer("Book Processing"):
    # Simulate some work
    books = []
    for i in range(1000000):
        books.append(i ** 2)
    time.sleep(0.5)

Output:

Starting Book Processing...
Book Processing took 0.5234 seconds

Practical Example: Temporary Directory Changer

from contextlib import contextmanager
import os

@contextmanager
def change_directory(new_dir):
    """Temporarily change working directory"""
    old_dir = os.getcwd()
    print(f"Changing directory: {old_dir} → {new_dir}")

    try:
        os.chdir(new_dir)
        yield
    finally:
        print(f"Restoring directory: {new_dir} → {old_dir}")
        os.chdir(old_dir)

# Example usage
print(f"Current directory: {os.getcwd()}")

with change_directory("/tmp"):
    print(f"Inside with block: {os.getcwd()}")
    # Do work in /tmp...

print(f"After with block: {os.getcwd()}")

Suppressing Exceptions with contextlib.suppress

from contextlib import suppress

# Instead of try/except for specific exceptions
books = ["Python Basics", "Python Advanced"]

# Without suppress
try:
    books.remove("Data Science 101")  # Not in list
except ValueError:
    pass

# With suppress (cleaner)
with suppress(ValueError):
    books.remove("Machine Learning")  # Not in list

print(books)

Redirecting Output with contextlib.redirect_stdout

from contextlib import redirect_stdout
import io

# Redirect print output to a string
output = io.StringIO()

with redirect_stdout(output):
    print("This goes to the string")
    print("So does this")

# Get the captured output
result = output.getvalue()
print(f"Captured: {result}")

Output:

Captured: This goes to the string
So does this

Nested Context Managers

You can use multiple context managers:

from contextlib import contextmanager

@contextmanager
def transaction(name):
    print(f"BEGIN {name}")
    try:
        yield
    finally:
        print(f"END {name}")

# Nested with statements
with transaction("Outer"):
    print("  Outer work")
    with transaction("Inner"):
        print("    Inner work")
    print("  More outer work")

Output:

BEGIN Outer
  Outer work
BEGIN Inner
    Inner work
END Inner
  More outer work
END Outer

You can also write them on one line:

with open('input.txt', 'r') as infile, open('output.txt', 'w') as outfile:
    content = infile.read()
    outfile.write(content.upper())

Practical Example: File Writer with Backup

Let’s create a context manager that creates a backup before writing:

from contextlib import contextmanager
import os
import shutil

@contextmanager
def safe_write(filename):
    """Context manager that backs up file before writing"""
    backup_name = f"{filename}.backup"

    # Create backup if file exists
    if os.path.exists(filename):
        print(f"Creating backup: {backup_name}")
        shutil.copy2(filename, backup_name)

    try:
        # Open file for writing
        with open(filename, 'w') as f:
            yield f

        # Success - remove backup
        if os.path.exists(backup_name):
            print(f"Removing backup: {backup_name}")
            os.remove(backup_name)

    except Exception as e:
        # Error - restore from backup
        print(f"Error occurred: {e}")
        if os.path.exists(backup_name):
            print(f"Restoring from backup: {backup_name}")
            shutil.move(backup_name, filename)
        raise

# Use it
with safe_write('books.txt') as f:
    f.write("Python Basics\n")
    f.write("Python Advanced\n")

Practical Example: Book Catalog Manager

Let’s build a comprehensive example combining everything we’ve learned:

from contextlib import contextmanager
import time

class BookCatalog:
    """Manage book catalog with context managers"""

    def __init__(self, name):
        self.name = name
        self.books = []
        self.locked = False

    @contextmanager
    def transaction(self):
        """Context manager for transactions"""
        print(f"[{self.name}] BEGIN TRANSACTION")
        original_books = self.books.copy()
        changes = []

        try:
            # Yield a helper object for making changes
            class TransactionHelper:
                def add_book(helper_self, book):
                    self.books.append(book)
                    changes.append(('ADD', book))
                    print(f"  Added: {book}")

                def remove_book(helper_self, book):
                    if book in self.books:
                        self.books.remove(book)
                        changes.append(('REMOVE', book))
                        print(f"  Removed: {book}")

            yield TransactionHelper()

            # Commit
            print(f"[{self.name}] COMMIT ({len(changes)} changes)")

        except Exception as e:
            # Rollback
            self.books = original_books
            print(f"[{self.name}] ROLLBACK ({len(changes)} changes)")
            raise

    @contextmanager
    def lock(self):
        """Context manager for locking catalog"""
        if self.locked:
            raise RuntimeError("Catalog is already locked")

        print(f"[{self.name}] LOCK acquired")
        self.locked = True

        try:
            yield
        finally:
            self.locked = False
            print(f"[{self.name}] LOCK released")

    def display(self):
        """Display catalog contents"""
        print(f"\n{self.name} Catalog ({len(self.books)} books):")
        for i, book in enumerate(self.books, 1):
            print(f"  {i}. {book}")

# Use the catalog with context managers
catalog = BookCatalog("Main Store")

# Successful transaction
print("=== Adding Books ===")
with catalog.lock():
    with catalog.transaction() as txn:
        txn.add_book("Python Basics")
        txn.add_book("Python Advanced")
        txn.add_book("Data Science 101")

catalog.display()

# Failed transaction (automatically rolls back)
print("\n=== Attempting Bad Transaction ===")
try:
    with catalog.transaction() as txn:
        txn.add_book("Machine Learning")
        txn.add_book("Deep Learning")
        raise ValueError("Inventory check failed!")
except ValueError as e:
    print(f"Transaction failed: {e}")

catalog.display()

# Removing a book
print("\n=== Removing a Book ===")
with catalog.transaction() as txn:
    txn.remove_book("Python Advanced")

catalog.display()

Output:

=== Adding Books ===
[Main Store] LOCK acquired
[Main Store] BEGIN TRANSACTION
  Added: Python Basics
  Added: Python Advanced
  Added: Data Science 101
[Main Store] COMMIT (3 changes)
[Main Store] LOCK released

Main Store Catalog (3 books):
  1. Python Basics
  2. Python Advanced
  3. Data Science 101

=== Attempting Bad Transaction ===
[Main Store] BEGIN TRANSACTION
  Added: Machine Learning
  Added: Deep Learning
[Main Store] ROLLBACK (2 changes)
Transaction failed: Inventory check failed!

Main Store Catalog (3 books):
  1. Python Basics
  2. Python Advanced
  3. Data Science 101

=== Removing a Book ===
[Main Store] BEGIN TRANSACTION
  Removed: Python Advanced
[Main Store] COMMIT (1 changes)

Main Store Catalog (2 books):
  1. Python Basics
  2. Data Science 101

Summary

In this lesson, we learned about context managers:

  • Context managers handle resource setup and cleanup automatically
  • The with statement uses __enter__ and __exit__ methods
  • __exit__ is always called, even if an exception occurs
  • Class-based: Implement __enter__ and __exit__ methods
  • Function-based: Use @contextmanager decorator with generators
  • contextlib module provides useful context manager utilities

Common uses:

  • File operations
  • Database transactions
  • Locks and synchronization
  • Temporary state changes
  • Timing and profiling
  • Exception suppression

Benefits:

  • Automatic cleanup (no forgotten close() calls)
  • Exception-safe (cleanup happens even with errors)
  • Cleaner code (no try/finally boilerplate)
  • Reusable resource management patterns

Best practices:

  • Use context managers for any resource that needs cleanup
  • Prefer @contextmanager for simple cases
  • Use classes for complex state management
  • Don’t suppress exceptions unless intentional (return False from __exit__)

In the next lesson, we’ll explore regular expressions—a powerful tool for text processing and pattern matching.

Prev
Lesson 8 - Generators and Memory-Efficient Processing
Next
Lesson 10 - Regular Expressions for Text Processing