Day 18: Context Managers — Beyond the with Statement

🐍 Day 18: Context Managers — Beyond the with Statement


1. Learning Objectives

By the end of Day 18, you will be able to:

  • Understand the purpose of context managers and the with statement
  • Implement a custom context manager using a class with __enter__ and __exit__
  • Use the contextlib module to create context managers with the @contextmanager decorator
  • Handle exceptions cleanly inside __exit__
  • Nest multiple context managers seamlessly
  • Apply context managers to real‑world resource management (files, database connections, locks, etc.)

2. Concept Explanation

2.1 Why Context Managers? – The Problem They Solve

Many resources need to be acquired and then released reliably – files must be closed, database connections returned to the pool, locks released. If you forget to release them, you get resource leaks, corrupt data, or deadlocks.

Without context managers (manual cleanup):

f = open("data.txt", "r")
try:
    content = f.read()
finally:
    f.close()

The finally block ensures the file is closed even if an error occurs. However, this pattern is repetitive and error‑prone. Python’s with statement and context managers encapsulate this pattern into a single, readable construct.


2.2 The with Statement – Clean and Reliable

We’ve used with for files since Day 9. Any object that implements the context manager protocol can be used with with:

with open("data.txt", "r") as f:
    content = f.read()
# File is automatically closed here – no matter what

Under the hood, two special methods are called:

  • __enter__(self) – called when entering the with block; returns the resource (often self).
  • __exit__(self, exc_type, exc_val, exc_tb) – called when leaving the block; handles cleanup and optionally suppresses exceptions.

2.3 Implementing a Context Manager with a Class

Let’s build a context manager that measures the execution time of a block of code:

import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        # We can return this instance itself (or any object)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time.perf_counter()
        self.elapsed = self.end - self.start
        print(f"Elapsed time: {self.elapsed:.4f} seconds")
        # Return False (or None) to propagate exceptions,
        # Return True to suppress them.
        return False   # Don't suppress exceptions

with Timer() as t:
    # simulate some work
    result = sum(range(1_000_000))
print("Done")

Output:

Elapsed time: 0.0234 seconds
Done

Key points:

  • __enter__ can return any object; it’s bound to the variable after as (here t).
  • __exit__ receives exception information. If no exception occurred, all three arguments are None.
  • If __exit__ returns True, any exception raised inside the block is suppressed. Usually you return False (or nothing, which defaults to False) so exceptions propagate normally.

2.4 Handling Exceptions Inside __exit__

A safe file‑opener that rolls back changes if an error occurs:

class SafeFileWriter:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            # An exception occurred – clean up or roll back
            print(f"An error occurred: {exc_val}. Rolling back changes.")
            self.file.close()
            # Optionally, delete the file or revert
            import os
            os.remove(self.filename)
        else:
            self.file.close()
        # Do not suppress the exception
        return False

with SafeFileWriter("important.txt", "w") as f:
    f.write("Critical data\n")
    # Simulate an error after writing
    raise ValueError("Something went wrong!")

The file will be closed and removed because an exception occurred.


2.5 The contextlib Module – @contextmanager Decorator

You can create a context manager more concisely using a generator function and the @contextmanager decorator from contextlib. The function must yield exactly once; code before yield is the __enter__, code after yield is the __exit__.

from contextlib import contextmanager

@contextmanager
def timer():
    import time
    start = time.perf_counter()
    yield                      # The resource (optional) can be yielded
    end = time.perf_counter()
    print(f"Elapsed: {end - start:.4f}s")

with timer():
    total = sum(range(1_000_000))

If you need to return a resource, yield it:

@contextmanager
def open_file(filename, mode):
    f = open(filename, mode)
    try:
        yield f                 # The file object is available as the "as" variable
    finally:
        f.close()

with open_file("test.txt", "r") as f:
    print(f.read())

⚠️ The finally block ensures cleanup even if an exception occurs inside the with block. Without it, the file would not be closed if an exception is raised.


2.6 Nested Context Managers

Multiple resources can be managed with a single with statement, separated by commas:

with open("input.txt") as infile, open("output.txt", "w") as outfile:
    for line in infile:
        outfile.write(line.upper())

This is equivalent to nested with blocks and guarantees all resources are closed properly.


2.7 Real‑World Use Cases

  • File handling: avoid resource leaks.
  • Database connections: ensure connections are returned to the pool.
  • Locks (threading): automatically acquire and release locks.
  • Temporary changes: temporarily change directory, environment variables, etc.
  • Profiling & benchmarking: measure execution time of code blocks.
  • Transaction management: roll back on error.

3. Code Examples

Example 1: Temporary Directory Cleaner

import os
import tempfile
from contextlib import contextmanager

@contextmanager
def temporary_directory():
    import shutil
    dirpath = tempfile.mkdtemp()
    try:
        yield dirpath
    finally:
        shutil.rmtree(dirpath)

with temporary_directory() as tmp:
    # Work inside the temporary directory
    with open(os.path.join(tmp, "data.txt"), "w") as f:
        f.write("temporary content")
    # The directory and its contents will be deleted automatically

Example 2: Database Transaction Simulator

class DatabaseSimulator:
    def __init__(self):
        self.data = {}

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

    def commit(self):
        print("Commit changes")

    def rollback(self):
        print("Rollback changes")

class Transaction:
    def __init__(self, db):
        self.db = db

    def __enter__(self):
        print("BEGIN TRANSACTION")
        return self.db

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            print("Error occurred – rolling back")
            self.db.rollback()
        else:
            print("Committing")
            self.db.commit()
        return False

db = DatabaseSimulator()

with Transaction(db) as conn:
    conn.execute("INSERT INTO users VALUES (1, 'Alice')")
    # If something goes wrong:
    # raise RuntimeError("Oops")

4. Hands-On Exercises

Exercise 1: Simple Timer Context Manager (Class)

Create a class Chronometer that implements __enter__ and __exit__. It should measure and print the elapsed time. Test it with a block of code.

Exercise 2: Timer with @contextmanager

Rewrite Exercise 1 using the @contextmanager decorator. Make sure it handles exceptions properly (the timer should still report even if an exception occurs).

Exercise 3: Safe File Writer with Rollback

Implement a context manager class SafeWrite that opens a file for writing, but if any exception occurs inside the block, it should close the file and delete it (rollback). If successful, just close the file.

Exercise 4: Suppress Specific Exception

Create a context manager suppress_exception(*exceptions) that suppresses (ignores) the given exception types. Use it to ignore FileNotFoundError when trying to open a missing file. (Hint: __exit__ can return True to suppress.)

Exercise 5: Nested Resources

Open two files with a single with statement: read from source.txt and write to destination.txt all lines that contain the word "Python". Use the nested syntax and proper error handling.


5. Applied Challenge Task 🏗️

Resource Pool Manager

Build a simple resource pool manager for database connections (simulated). The pool has a fixed number of connections.

Requirements:

  1. Class Connection: just a dummy that prints "Connected" when created and "Disconnected" when released.

  2. Class ConnectionPool:

    • Manages a list of available connections (maximum max_connections).
    • Has a method acquire() that returns a connection (waits if none available – but for simplicity, just raise RuntimeError if exhausted).
    • Has a method release(conn) that returns the connection to the pool.
  3. Context Manager pooled_connection(pool):

    • Acquires a connection from the pool on enter.
    • Releases it on exit.
    • If an exception occurs inside the with block, still release the connection.
  4. Demonstrate:

    • Create a pool with 2 connections.
    • Use two nested with blocks to acquire both connections (should work).
    • Try to acquire a third connection (should raise error).
    • Show that after exiting a with block, the connection is returned.

Stretch goals:

  • Use @contextmanager for the pooled connection.
  • Simulate a timeout waiting for a free connection.
  • Log all acquire/release events to a file.

6. Brief Review Summary

ConceptKey Points
Context manager protocol__enter__ and __exit__ methods
with statementAutomates setup and teardown; ensures cleanup
__exit__ argumentsexc_type, exc_val, exc_tb – exception info
Suppressing exceptionsReturn True from __exit__
@contextmanagerGenerator‑based context managers from contextlib
NestingMultiple managers in one with separated by commas
Use casesFiles, locks, DB connections, temporary changes

7. Preview of Next Topic — Day 19

Tomorrow we’ll step into real‑world data exchange:

  • Working with APIs – making HTTP requests
  • JSON handling – parsing and generating JSON data
  • The requests library and json module
  • Building a simple API client

🎯 Your Action Items for Day 18:

  1. ✅ Complete all 5 exercises
  2. ✅ Build the Resource Pool Manager challenge
  3. ✅ Convert an existing file‑based code to use a custom context manager
  4. ✅ Experiment with contextlib.suppress and compare with a custom suppressor

Comments

Popular posts from this blog

Day 1: Welcome to Python — Your Journey Begins

Python: Your Gateway to Coding Adventures

Day 11: List Comprehensions & Lambda Functions

Earn From the Comfort of Your Home

Build the skills to work comfortably from home, on your own terms.

Get Started Today