Context Managers

Ensure proper resource cleanup with context managers.

Context managers in Python are used to manage resources, ensuring they are properly acquired and released. They are typically used with the with statement and help prevent resource leaks.

Why Use Context Managers?

Without context managers (error-prone):

1
file = open('file.txt', 'r')
2
try:
3
    content = file.read()
4
finally:
5
    file.close()  # Must remember to close!

With context managers (pythonic):

1
with open('file.txt', 'r') as file:
2
    content = file.read()
3
# File is automatically closed after the block

Creating Custom Context Managers

Class-Based Context Manager

A context manager needs to implement __enter__ and __exit__ methods:

1
class DatabaseConnection:
2
    """Custom context manager for database connections"""
3
    def __init__(self, connection_string: str):
4
        self.connection_string = connection_string
5
        self.connection = None
6

7
    def __enter__(self):
8
        """Called when entering the 'with' block"""
9
        print("Connecting to database...")
10
        self.connection = self.connect_to_database()
11
        return self.connection
12

13
    def __exit__(self, exc_type, exc_value, traceback):
14
        """Called when exiting the 'with' block"""
15
        print("Closing database connection...")
16
        if self.connection:
17
            self.connection.close()
18
        return False  # Don't suppress exceptions
19

20
    def connect_to_database(self):
21
        # Simulate database connection
22
        return {"status": "connected"}
23

24
# Usage
25
with DatabaseConnection("postgresql://localhost/db") as db:
26
    print("Using database:", db)
27
    # Database operations here
28
# Connection automatically closed

Function-Based Context Manager (using `contextlib`)

The contextlib module provides a decorator for creating context managers:

1
from contextlib import contextmanager
2

3
@contextmanager
4
def database_connection(connection_string: str):
5
    """Function-based context manager"""
6
    print("Connecting to database...")
7
    connection = connect_to_database(connection_string)
8
    try:
9
        yield connection  # Value returned to 'as' variable
10
    finally:
11
        print("Closing database connection...")
12
        connection.close()
13

14
def connect_to_database(connection_string: str):
15
    return {"status": "connected"}
16

17
# Usage
18
with database_connection("postgresql://localhost/db") as db:
19
    print("Using database:", db)

Real-World Examples

Timer Context Manager

1
import time
2
from contextlib import contextmanager
3

4
@contextmanager
5
def timer():
6
    """Context manager to measure execution time"""
7
    start = time.time()
8
    try:
9
        yield
10
    finally:
11
        end = time.time()
12
        print(f"Execution time: {end - start:.2f} seconds")
13

14
# Usage
15
with timer():
16
    # Your code here
17
    time.sleep(1)
18
    print("Task completed")
19
# Output: Execution time: 1.00 seconds

Temporary Directory Context Manager

1
import tempfile
2
import shutil
3
from contextlib import contextmanager
4

5
@contextmanager
6
def temporary_directory():
7
    """Create and clean up a temporary directory"""
8
    temp_dir = tempfile.mkdtemp()
9
    try:
10
        yield temp_dir
11
    finally:
12
        shutil.rmtree(temp_dir)  # Clean up
13

14
# Usage
15
with temporary_directory() as temp_dir:
16
    # Work with temp_dir
17
    print(f"Working in {temp_dir}")
18
# Directory automatically deleted

Lock Context Manager

1
import threading
2
from contextlib import contextmanager
3

4
@contextmanager
5
def acquire_lock(lock: threading.Lock):
6
    """Acquire and release a lock"""
7
    lock.acquire()
8
    try:
9
        yield
10
    finally:
11
        lock.release()
12

13
# Usage
14
lock = threading.Lock()
15
with acquire_lock(lock):
16
    # Critical section
17
    print("Thread-safe code here")
18
# Lock automatically released

Understanding `exit` Parameters

The __exit__ method receives three parameters about any exception that occurred:

1
class ErrorLogger:
2
    """Context manager that logs exceptions"""
3
    def __enter__(self):
4
        return self
5

6
    def __exit__(self, exc_type, exc_value, traceback):
7
        if exc_type is not None:
8
            print(f"Exception occurred: {exc_type.__name__}: {exc_value}")
9
            # Return False to propagate the exception
10
            # Return True to suppress it
11
        return False
12

13
# Usage
14
with ErrorLogger():
15
    raise ValueError("Something went wrong")
16
# Exception is logged but still raised

Multiple Context Managers

You can use multiple context managers in one statement:

1
from contextlib import contextmanager
2
import time
3

4
@contextmanager
5
def timer():
6
    start = time.time()
7
    try:
8
        yield
9
    finally:
10
        print(f"Time: {time.time() - start:.2f}s")
11

12
@contextmanager
13
def logger(name: str):
14
    print(f"Starting {name}...")
15
    try:
16
        yield
17
    finally:
18
        print(f"Finished {name}")
19

20
# Multiple context managers
21
with timer(), logger("operation"):
22
    time.sleep(1)
23
    print("Doing work...")