Context Managers and Python's with Statement

The try … finally Approach

Working with files is probably the most common example of resource management in programming. In Python, you can use a try … finally statement to handle opening and closing files properly:

# Safely open the file
file = open("hello.txt", "w")

try:
    file.write("Hello, World!")
finally:
    # Make sure to close the file after using it
    file.close()

In this example, you need to safely open the file hello.txt, which you can do by wrapping the call to open() in a try … except statement. Later, when you try to write to file, the finally clause will guarantee that file is properly closed, even if an exception occurs during the call to .write() in the try clause. You can use this pattern to handle setup and teardown logic when you’re managing external resources in Python.

The try block in the above example can potentially raise exceptions, such as AttributeError or NameError. You can handle those exceptions in an except clause like this:

# Safely open the file
file = open("hello.txt", "w")

try:
    file.write("Hello, World!")
except Exception as e:
    print(f"An error occurred while writing to the file: {e}")
finally:
    # Make sure to close the file after using it
    file.close()

In this example, you catch any potential exceptions that can occur while writing to the file. In real-life situations, you should use a specific exception type instead of the general Exception to prevent unknown errors from passing silently.

The with Statement Approach

The Python with statement creates a runtime context that allows you to run a group of statements under the control of a context manager. PEP 343 added the with statement to make it possible to factor out standard use cases of the try … finally statement.

Compared to traditional try … finally constructs, the with statement can make your code clearer, safer, and reusable. Many classes in the standard library support the with statement. A classic example of this is open(), which allows you to work with file objects using with.

To write a with statement, you need to use the following general syntax:

with expression as target_var:
    do_something(target_var)

The context manager object results from evaluating the expression after with. In other words, expression must return an object that implements the context management protocol. This protocol consists of two special methods:

1. .__enter__() is called by the with statement to enter the runtime context.
2. .__exit__() is called when the execution leaves the with code block.

The as specifier is optional. If you provide a target_var with as, then the return value of calling .__enter__() on the context manager object is bound to that variable.

Here’s how the with statement proceeds when Python runs into it:

1. Call expression to obtain a context manager.
2. Store the context manager’s .__enter__() and .__exit__() methods for later use.
3. Call .__enter__() on the context manager and bind its return value to target_var if provided.
4. Execute the with code block.
5. Call .__exit__() on the context manager when the with code block finishes.

In this case, .__enter__(), typically provides the setup code. The with statement is a compound statement that starts a code block, like a conditional statement or a for loop. Inside this code block, you can run several statements. Typically, you use the with code block to manipulate target_var if applicable.

Once the with code block finishes, .__exit__() gets called. This method typically provides the teardown logic or cleanup code, such as calling .close() on an open file object. That’s why the with statement is so useful. It makes properly acquiring and releasing resources a breeze.

Here’s how to open your hello.txt file for writing using the with statement:

with open("hello.txt", mode="w") as file:
    file.write("Hello, World!")

When you run this with statement, open() returns an io.TextIOBase object. This object is also a context manager, so the with statement calls .__enter__() and assigns its return value to file. Then you can manipulate the file inside the with code block. When the block ends, .__exit__() automatically gets called and closes the file for you, even if an exception is raised inside the with block.

This with construct is shorter than its try … finally alternative, but it’s also less general, as you already saw. You can only use the with statement with objects that support the context management protocol, whereas try … finally allows you to perform cleanup actions for arbitrary objects without the need for supporting the context management protocol.

In Python 3.1 and later, the with statement supports multiple context managers. You can supply any number of context managers separated by commas:

with A() as a, B() as b:
    pass

This works like nested with statements but without nesting. This might be useful when you need to open two files at a time, the first for reading and the second for writing:

with open("input.txt") as in_file, open("output.txt", "w") as out_file:
    # Read content from input.txt
    # Transform the content
    # Write the transformed content to output.txt
    pass

In this example, you can add code for reading and transforming the content of input.txt. Then you write the final result to output.txt in the same code block.

Using multiple context managers in a single with has a drawback, though. If you use this feature, then you’ll probably break your line length limit. To work around this, you need to use backslashes (\) for line continuation, so you might end up with an ugly final result.

The with statement can make the code that deals with system resources more readable, reusable, and concise, not to mention safer. It helps avoid bugs and leaks by making it almost impossible to forget cleaning up, closing, and releasing a resource after you’re done with it.

Using with allows you to abstract away most of the resource handling logic. Instead of having to write an explicit try … finally statement with setup and teardown code each time, with takes care of that for you and avoids repetition.

Working With Files

So far, you’ve used open() to provide a context manager and manipulate files in a with construct. Opening files using the with statement is generally recommended because it ensures that open file descriptors are automatically closed after the flow of execution leaves the with code block.

As you saw before, the most common way to open a file using with is through the built-in open():

with open("hello.txt", mode="w") as file:
    file.write("Hello, World!")

In this case, since the context manager closes the file after leaving the with code block, a common mistake might be the following:

>>> file = open("hello.txt", mode="w")

>>> with file:
...     file.write("Hello, World!")
...
13

>>> with file:
...     file.write("Welcome to Real Python!")
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: I/O operation on closed file.

The first with successfully writes "Hello, World!" into hello.txt. Note that .write() returns the number of bytes written into the file, 13. When you try to run a second with, however, you get a ValueError because your file is already closed.

Another way to use the with statement to open and manage files is by using pathlib.Path.open():

>>> import pathlib

>>> file_path = pathlib.Path("hello.txt")

>>> with file_path.open("w") as file:
...     file.write("Hello, World!")
...
13

Path is a class that represents concrete paths to physical files in your computer. Calling .open() on a Path object that points to a physical file opens it just like open() would do. So, Path.open() works similarly to open(), but the file path is automatically provided by the Path object you call the method on.

Since pathlib provides an elegant, straightforward, and Pythonic way to manipulate file system paths, you should consider using Path.open() in your with statements as a best practice in Python.

Finally, whenever you load an external file, your program should check for possible issues, such as a missing file, writing and reading access, and so on. Here’s a general pattern that you should consider using when you’re working with files:

import pathlib
import logging

file_path = pathlib.Path("hello.txt")

try:
    with file_path.open(mode="w") as file:
        file.write("Hello, World!")
except OSError as error:
    logging.error("Writing to file %s failed due to: %s", file_path, error)

In this example, you wrap the with statement in a try … except statement. If an OSError occurs during the execution of with, then you use logging to log the error with a user-friendly and descriptive message.

Traversing Directories

The os module provides a function called scandir(), which returns an iterator over os.DirEntry objects corresponding to the entries in a given directory. This function is specially designed to provide optimal performance when you’re traversing a directory structure.

A call to scandir() with the path to a given directory as an argument returns an iterator that supports the context management protocol:

>>> import os

>>> with os.scandir(".") as entries:
...     for entry in entries:
...         print(entry.name, "->", entry.stat().st_size, "bytes")
...
Documents -> 4096 bytes
Videos -> 12288 bytes
Desktop -> 4096 bytes
DevSpace -> 4096 bytes
.profile -> 807 bytes
Templates -> 4096 bytes
Pictures -> 12288 bytes
Public -> 4096 bytes
Downloads -> 4096 bytes

In this example, you write a with statement with os.scandir() as the context manager supplier. Then you iterate over the entries in the selected directory (".") and print their name and size on the screen. In this case, .__exit__() calls scandir.close() to close the iterator and release the acquired resources. Note that if you run this on your machine, you’ll get a different output depending on the content of your current directory.

Performing High-Precision Calculations

Unlike built-in floating-point numbers, the decimal module provides a way to adjust the precision to use in a given calculation that involves Decimal numbers. The precision defaults to 28 places, but you can change it to meet your problem requirements. A quick way to perform calculations with a custom precision is using localcontext() from decimal:

>>> from decimal import Decimal, localcontext

>>> with localcontext() as ctx:
...     ctx.prec = 42
...     Decimal("1") / Decimal("42")
...
Decimal('0.0238095238095238095238095238095238095238095')

>>> Decimal("1") / Decimal("42")
Decimal('0.02380952380952380952380952381')

Here, localcontext() provides a context manager that creates a local decimal context and allows you to perform calculations using a custom precision. In the with code block, you need to set .prec to the new precision you want to use, which is 42 places in the example above. When the with code block finishes, the precision is reset back to its default value, 28 places.

Handling Locks in Multithreaded Programs

Another good example of using the with statement effectively in the Python standard library is threading.Lock. This class provides a primitive lock to prevent multiple threads from modifying a shared resource at the same time in a multithreaded application.

You can use a Lock object as the context manager in a with statement to automatically acquire and release a given lock. For example, say you need to protect the balance of a bank account:

import threading

balance_lock = threading.Lock()

# Use the try ... finally pattern
balance_lock.acquire()
try:
    # Update the account balance here ...
finally:
    balance_lock.release()

# Use the with pattern
with balance_lock:
    # Update the account balance here ...

The with statement in the second example automatically acquires and releases a lock when the flow of execution enters and leaves the statement. This way, you can focus on what really matters in your code and forget about those repetitive operations.

In this example, the lock in the with statement creates a protected region known as the critical section, which prevents concurrent access to the account balance.

Testing for Exceptions With pytest

So far, you’ve coded a few examples using context managers that are available in the Python standard library. However, several third-party libraries include objects that support the context management protocol.

Say you’re testing your code with pytest. Some of your functions and code blocks raise exceptions under certain situations, and you want to test those cases. To do that, you can use pytest.raises(). This function allows you to assert that a code block or a function call raises a given exception.

Since pytest.raises() provides a context manager, you can use it in a with statement like this:

>>> import pytest

>>> 1 / 0
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ZeroDivisionError: division by zero

>>> with pytest.raises(ZeroDivisionError):
...     1 / 0
...

>>> favorites = {"fruit": "apple", "pet": "dog"}
>>> favorites["car"]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
KeyError: 'car'

>>> with pytest.raises(KeyError):
...     favorites["car"]
...

In the first example, you use pytest.raises() to capture the ZeroDivisionError that the expression 1 / 0 raises. The second example uses the function to capture the KeyError that is raised when you access a key that doesn’t exist in a given dictionary.

If your function or code block doesn’t raise the expected exception, then pytest.raises() raises a failure exception:

>>> import pytest

>>> with pytest.raises(ZeroDivisionError):
...     4 / 2
...
2.0
Traceback (most recent call last):
  ...
Failed: DID NOT RAISE <class 'ZeroDivisionError'>

Another cool feature of pytest.raises() is that you can specify a target variable to inspect the raised exception. For example, if you want to verify the error message, then you can do something like this:

>>> with pytest.raises(ZeroDivisionError) as exc:
...     1 / 0
...
>>> assert str(exc.value) == "division by zero"

You can use all these pytest.raises() features to capture the exceptions you raise from your functions and code block. This is a cool and useful tool that you can incorporate into your current testing strategy.

Writing a Sample Class-Based Context Manager

Here’s a sample class-based context manager that implements both methods, .__enter__() and .__exit__(). It also shows how Python calls them in a with construct:

>>> class HelloContextManager:
...     def __enter__(self):
...         print("Entering the context...")
...         return "Hello, World!"
...     def __exit__(self, exc_type, exc_value, exc_tb):
...         print("Leaving the context...")
...         print(exc_type, exc_value, exc_tb, sep="\n")
...

>>> with HelloContextManager() as hello:
...     print(hello)
...
Entering the context...
Hello, World!
Leaving the context...
None
None
None

HelloContextManager implements both .__enter__() and .__exit__(). In .__enter__(), you first print a message to signal that the flow of execution is entering a new context. Then you return the "Hello, World!" string. In .__exit__(), you print a message to signal that the flow of execution is leaving the context. You also print the content of its three arguments.

When the with statement runs, Python creates a new instance of HelloContextManager and calls its .__enter__() method. You know this because you get Entering the context... printed on the screen.

>>> with HelloContextManager as hello:
...     print(hello)
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: __enter__

Then Python runs the with code block, which prints hello to the screen. Note that hello holds the return value of .__enter__().

When the flow of execution exits the with code block, Python calls .__exit__(). You know that because you get Leaving the context... printed on your screen. The final line in the output confirms that the three arguments to .__exit__() are set to None.

Now, what happens if an exception occurs during the execution of the with block? Go ahead and write the following with statement:

>>> with HelloContextManager() as hello:
...     print(hello)
...     hello[100]
...
Entering the context...
Hello, World!
Leaving the context...
<class 'IndexError'>
string index out of range
<traceback object at 0x7f0cebcdd080>
Traceback (most recent call last):
  File "<stdin>", line 3, in <module>
IndexError: string index out of range

In this case, you try to retrieve the value at index 100 in the string "Hello, World!". This raises an IndexError, and the arguments to .__exit__() are set to the following:

• exc_type is the exception class, IndexError.
• exc_value is the exception instance.
• exc_tb is the traceback object.

This behavior is quite useful when you want to encapsulate the exception handling in your context managers.

Handling Exceptions in a Context Manager

As an example of encapsulating exception handling in a context manager, say you expect IndexError to be the most common exception when you’re working with HelloContextManager. You might want to handle that exception in the context manager so you don’t have to repeat the exception-handling code in every with code block. In that case, you can do something like this:

# exc_handling.py

class HelloContextManager:
    def __enter__(self):
        print("Entering the context...")
        return "Hello, World!"

    def __exit__(self, exc_type, exc_value, exc_tb):
        print("Leaving the context...")
        if isinstance(exc_value, IndexError):
            # Handle IndexError here...
            print(f"An exception occurred in your with block: {exc_type}")
            print(f"Exception message: {exc_value}")
            return True

with HelloContextManager() as hello:
    print(hello)
    hello[100]

print("Continue normally from here...")

In .__exit__(), you check if exc_value is an instance of IndexError. If so, then you print a couple of informative messages and finally return with True. Returning a truthy value makes it possible to swallow the exception and continue the normal execution after the with code block.

In this example, if no IndexError occurs, then the method returns None and the exception propagates out. However, if you want to be more explicit, then you can return False from outside the if block.

If you run exc_handling.py from your command line, then you get the following output:

$ python exc_handling.py
Entering the context...
Hello, World!
Leaving the context...
An exception occurred in your with block: <class 'IndexError'>
Exception message: string index out of range
Continue normally from here...

HelloContextManager is now able to handle IndexError exceptions that occur in the with code block. Since you return True when an IndexError occurs, the flow of execution continues in the next line, right after exiting the with code block.

Opening Files for Writing: First Version

Now that you know how to implement the context management protocol, you can get a sense of what this would look like by coding a practical example. Here’s how you can take advantage of open() to create a context manager that opens files for writing:

# writable.py

class WritableFile:
    def __init__(self, file_path):
        self.file_path = file_path

    def __enter__(self):
        self.file_obj = open(self.file_path, mode="w")
        return self.file_obj

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file_obj:
            self.file_obj.close()

WritableFile implements the context management protocol and supports the with statement, just like the original open() does, but it always opens the file for writing using the "w" mode. Here’s how you can use your new context manager:

>>> from writable import WritableFile

>>> with WritableFile("hello.txt") as file:
...    file.write("Hello, World!")
...

After running this code, your hello.txt file contains the "Hello, World!" string. As an exercise, you can write a complementary context manager that opens files for reading, but using pathlib functionalities. Go ahead and give it a shot!

Redirecting the Standard Output

A subtle detail to consider when you’re writing your own context managers is that sometimes you don’t have a useful object to return from .__enter__() and therefore to assign to the with target variable. In those cases, you can return None explicitly or you can just rely on Python’s implicit return value, which is None as well.

For example, say you need to temporarily redirect the standard output, sys.stdout, to a given file on your disk. To do this, you can create a context manager like this:

# redirect.py

import sys

class RedirectedStdout:
    def __init__(self, new_output):
        self.new_output = new_output

    def __enter__(self):
        self.saved_output = sys.stdout
        sys.stdout = self.new_output

    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.stdout = self.saved_output

This context manager takes a file object through its constructor. In .__enter__(), you reassign the standard output, sys.stdout, to an instance attribute to avoid losing the reference to it. Then you reassign the standard output to point to the file on your disk. In .__exit__(), you just restore the standard output to its original value.

To use RedirectedStdout, you can do something like this:

>>> from redirect import RedirectedStdout

>>> with open("hello.txt", "w") as file:
...     with RedirectedStdout(file):
...         print("Hello, World!")
...     print("Back to the standard output...")
...
Back to the standard output...

The outer with statement in this example provides the file object that you’re going to use as your new output, hello.txt. The inner with temporarily redirects the standard output to hello.txt, so the first call to print() writes directly to that file instead of printing "Hello, World!" on your screen. Note that when you leave the inner with code block, the standard output goes back to its original value.

RedirectedStdout is a quick example of a context manager that doesn’t have a useful value to return from .__enter__(). However, if you’re only redirecting the print() output, you can get the same functionality without the need for coding a context manager. You just need to provide a file argument to print() like this:

>>> with open("hello.txt", "w") as file:
...     print("Hello, World!", file=file)
...

In this examples, print() takes your hello.txt file as an argument. This causes print() to write directly into the physical file on your disk instead of printing "Hello, World!" to your screen.

Measuring Execution Time

Just like every other class, a context manager can encapsulate some internal state. The following example shows how to create a stateful context manager to measure the execution time of a given code block or function:

# timing.py

from time import perf_counter

class Timer:
    def __enter__(self):
        self.start = perf_counter()
        self.end = 0.0
        return lambda: self.end - self.start

    def __exit__(self, *args):
        self.end = perf_counter()

When you use Timer in a with statement, .__enter__() gets called. This method uses time.perf_counter() to get the time at the beginning of the with code block and stores it in .start. It also initializes .end and returns a lambda function that computes a time delta. In this case, .start holds the initial state or time measurement.

Once the with block ends, .__exit__() gets called. The method gets the time at the end of the block and updates the value of .end so that the lambda function can compute the time required to run the with code block.

Here’s how you can use this context manager in your code:

>>> from time import sleep
>>> from timing import Timer

>>> with Timer() as timer:
...     # Time-consuming code goes here...
...     sleep(0.5)
...

>>> timer()
0.5005456680000862

With Timer, you can measure the execution time of any piece of code. In this example, timer holds an instance of the lambda function that computes the time delta, so you need to call timer() to get the final result.

Opening Files for Writing: Second Version

You can use the @contextmanager to reimplement your WritableFile context manager. Here’s what rewriting it with this technique looks like:

>>> from contextlib import contextmanager

>>> @contextmanager
... def writable_file(file_path):
...     file = open(file_path, mode="w")
...     try:
...         yield file
...     finally:
...         file.close()
...

>>> with writable_file("hello.txt") as file:
...     file.write("Hello, World!")
...

In this case, writable_file() is a generator function that opens file for writing. Then it temporarily suspends its own execution and yields the resource so with can bind it to its target variable. When the flow of execution leaves the with code block, the function continues to execute and closes file correctly.

Mocking the Time

As a final example of how to create custom context managers with @contextmanager, say you’re testing a piece of code that works with time measurements. The code uses time.time() to get the current time measurement and do some further computations. Since time measurements vary, you decide to mock time.time() so you can test your code.

Here’s a function-based context manager that can help you do that:

>>> from contextlib import contextmanager
>>> from time import time

>>> @contextmanager
... def mock_time():
...     global time
...     saved_time = time
...     time = lambda: 42
...     yield
...     time = saved_time
...

>>> with mock_time():
...     print(f"Mocked time: {time()}")
...
Mocked time: 42

>>> # Back to normal time
>>> time()
1616075222.4410584

Inside mock_time(), you use a global statement to signal that you’re going to modify the global name time. Then you save the original time() function object in saved_time so you can safely restore it later. The next step is to monkey patch time() using a lambda function that always returns the same value, 42.

The bare yield statement specifies that this context manager doesn’t have a useful object to send back to the with target variable for later use. After yield, you reset the global time to its original content.

When the execution enters the with block, any calls to time() return 42. Once you leave the with code block, calls to time() return the expected current time. That’s it! Now you can test your time-related code.