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.