Python's with Statement: Manage External Resources Safely

The try … finally Construct

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

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

try:
    file.write("Hello, World!")
finally:
    file.close()

In this example, you open the hello.txt file using open(). To write some text into the file, you wrap the call to .write() in a try statement with a finally clause. This clause guarantees that the file is properly closed by calling .close(), even if an exception occurs during the call to .write() in the try clause. Remember that the finally clause always runs.

When managing external resources in Python, you can use the construct in the previous example to handle setup and teardown logic. The setup logic might include opening the file and writing content to it, while the teardown logic might consist of closing the file to release the acquired resources.

The try block in the example above can potentially raise exceptions, such as AttributeError or NameError. You can handle these exceptions with an except clause, as shown in the following example:

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:
    file.close()

In this example, you use the Exception class to catch exceptions that can occur while writing to the file. Then, you print an error message to inform the user. The finally clause always runs so that the file is closed correctly.

Even though the try … finally construct works correctly, Python has a more compact and elegant solution for you: the with statement.

The with Statement

The Python with statement creates a runtime context that allows you to execute a code block under the control of a context manager. PEP 343 added the with statement to make it possible to factor out common use cases of the try … finally construct.

Compared to the traditional try … finally construct, the with statement makes your code clearer, safer, and more readable. Many classes and objects in the standard library support the with statement by implementing the context manager protocol (special methods). For example, when you call open(), you get a file object that supports this protocol and, therefore, the with statement.

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

with expression [as obj]:
    <block>

The context manager object results from evaluating the expression after the with keyword. 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 an obj variable 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 works internally:

1. Execute expression to obtain a context manager object.
2. Call .__enter__() on the context manager and bind its return value to obj if provided.
3. Execute the with code block.
4. Call .__exit__() on the context manager when the with code block finishes.

The .__enter__() method typically provides the setup logic. The with statement is a compound statement that starts a code block, like other compound statements, such as conditional statements and for loops.

Inside this code block, you can run one or more statements. Typically, you use the with code block to manipulate obj if applicable.

Once the with code block finishes, Python automatically calls .__exit__(). 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", encoding="utf-8") as file:
    file.write("Hello, World!")

When you run this with statement, open() returns an io.TextIOBase object. This object supports the context manager protocol, 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, Python automatically calls .__exit__(), which closes the file for you, even if an exception occurs in the with code block.

This with construct is shorter than its try … finally equivalent construct, but it’s also less generic because you can only use the with statement with objects that support the context management protocol. In contrast, try … finally allows you to perform cleanup actions for any object, even if it doesn’t implement the context management protocol.

The with statement also supports multiple context managers. You can supply any number of context managers separated by commas:

with A() as a, B() as b:
    <block>

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

with open("hello.txt") as in_file, open("output.txt", "w") as out_file:
    for line in in_file:
        out_file.write(line.upper())

In this example, you read and transform the content of hello.txt line by line. Then, you write each processed line to output.txt within the same code block.

The with statement can make the code that deals with system resources more readable, concise, and safer. It helps avoid resource leaks by making it almost impossible to forget to clean up, close, and release resources after you’re done with them.

Using with allows you to abstract away most of the resource handling logic. Instead of using explicit try … finally constructs with setup and teardown logic, you can pack this logic into a context manager and handle it using the with statement to avoid repetition.

Working With Files

So far, you’ve used open() to manipulate files in a with construct. Managing files this way is generally recommended because it ensures that opened file descriptors are automatically closed after the flow of execution leaves the with code block.

Again, a common way to open a file in a with statement is through the open() function:

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

In this example, the context manager closes the file after leaving the with code block. A common mistake you might encounter is shown below:

>>> file = open("hello.txt", mode="w", encoding="utf-8")

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

>>> with file:
...     file.write("Welcome to Real Python!")
...
Traceback (most recent call last):
    ...
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 bytes in this example. When you try to run a second with, you get a ValueError because the file is already closed.

Another way to use the with statement to open and manage files is by using the .open() method on a Path object:

>>> import pathlib

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

>>> with file_path.open("w", encoding="utf-8") as file:
...     file.write("Hello, World!")
...
13

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

The pathlib package 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 pattern that you should consider using when 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)
...
ERROR:root:Writing to file /hello.txt failed due to: [Errno 13]
⮑ Permission denied: '/hello.txt'

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

Traversing Directories

The os module provides a function called scandir() that 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 a directory path 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 working directory represented by "." and print their names and sizes 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 example on your machine, you’ll get a different output depending on the content of your target directory.

Performing High-Precision Calculations

The decimal module provides a handy way to tweak 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 specific requirements.

A quick way to perform calculations with a custom precision is using localcontext() from decimal:

>>> from decimal import Decimal, localcontext

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

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

Here, localcontext() returns a context manager that creates a context allowing you to perform calculations using a custom precision. In the with code block, you have a precision of 42 places. When the with code block finishes, the precision is back to its default value of 28 places.

Handling Locks in Multithreaded Programs

The threading.Lock class in the Python standard library also supports the with statement. 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 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 also 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(), which allows you to assert that a block of code or a function call raises a specific exception.

The pytest.raises() function returns a context manager, so you can use that object in a with statement:

>>> import pytest

>>> 1 / 0
Traceback (most recent call last):
    ...
ZeroDivisionError: division by zero

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

>>> favorites = {"fruit": "apple", "pet": "dog"}
>>> favorites["car"]
Traceback (most recent call last):
    ...
KeyError: 'car'

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

In the first example, you use pytest.raises() to capture the ZeroDivisionError exception that the expression 1 / 0 raises. The second example uses this function to capture the KeyError that’s 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:

>>> 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 blocks. This is a cool and useful tool that you can incorporate into your current testing strategy.

Writing a Demo Class-Based Context Manager

When the with statement executes, Python calls .__enter__() on the context manager object to set up the new runtime context. If you provide a target variable with the as specifier, then the return value of .__enter__() is assigned to that variable. When the flow of execution leaves the context, .__exit__() is called.

Below is a demo class-based context manager that shows how Python calls the .__enter__() and .__exit__() methods 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(f"{exc_type  = }")
...         print(f"{exc_value = }")
...         print(f"{exc_tb    = }")
...
>>> with HelloContextManager() as hello:
...     print(hello)
...
Entering the context...
Hello, World!
Leaving the context...
exc_type  = None
exc_value = None
exc_tb    = None

In .__enter__(), you print a message to indicate 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 execution flow is leaving the context. You also print the content of the 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. Next, Python runs the with code block, which prints hello to the screen. Note that hello holds the return value of .__enter__(), which is "Hello, World!" in this example.

When the execution flow exits the with code block, Python calls .__exit__(). You know that because you get Leaving the context... printed on your screen. If no exception occurs in the with code block, then the three arguments to .__exit__() are set to None. You confirm this behavior in the final output lines.

Now, what if an exception occurs during the execution of the with block? To find out, go ahead and run the following with statement:

>>> with HelloContextManager() as hello:
...     print(hello)
...     hello[100]
...
Entering the context...
Hello, World!
Leaving the context...
exc_type  = <class 'IndexError'>
exc_value = IndexError('string index out of range')
exc_tb    = <traceback object at 0x7f0cebcdd080>
Traceback (most recent call last):
  ...
IndexError: string index out of range

In this example, you try to retrieve the value at index 100 in the string "Hello, World!". This action raises an IndexError exception because the string doesn’t have 100 characters. Therefore, the arguments to .__exit__() are set as follows:

• exc_type: the exception class, IndexError.
• exc_value: the concrete exception object with a descriptive message.
• exc_tb: the traceback object, <traceback object at 0x7f0cebcdd080>.

This behavior is especially useful when you want your context manager to handle exceptions internally.

Handling Exceptions Within Context Managers

If the .__exit__() method returns True, then any exception that occurs in the with block is swallowed and the execution continues at the next statement after with. If .__exit__() returns False, then exceptions are propagated out of the context. This is also the default behavior when the method doesn’t return anything explicitly.

You can take advantage of these behaviors to encapsulate exception handling inside the context manager or to propagate the exceptions as needed.

As an example of encapsulating exception handling in a context manager, say you expect IndexError to be a possible 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 situation, you can do something like the following:

>>> 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]
...
Entering the context...
Hello, World!
Leaving the context...
An exception occurred in your with block: <class 'IndexError'>
Exception message: string index out of range

In .__exit__(), you check if exc_value is an instance of IndexError using the built-in isinstance() function. If so, you print a couple of informative messages and finally return with True.

Again, returning a True 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 implicitly returns None, and the exception propagates out of the context. This is the desired behavior because it avoids hiding unexpected exceptions. If you want to be more explicit, then you can return False outside the if block.

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

Opening Files for Writing

Now that you know how to implement the context management protocol, you can code a more realistic example. For instance, below is how you can create a context manager that opens files for writing directly:

>>> 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, *_):
...         if self.file_obj:
...             self.file_obj.close()
...

WritableFile allows you to open the file for writing using the "w" mode of open(). You do this in the .__enter__() method. In the .__exit__() method, you close the file to release the acquired resources. Note how you use the *args syntax to maintain compatibility with the context management protocol. Since you don’t need these arguments, you use a single underscore to suppress a potential linter warning about unused variables.

Here’s how you can use your context manager:

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

After running this code, you’ll have a hello.txt file containing 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 try!

Redirecting the Standard Output to a File

A subtle detail to consider when you’re implementing the context manager protocol is that sometimes you don’t have a useful object to return from .__enter__()—and therefore nothing meaningful to assign to the target variable in the with header. In those cases, you can explicitly return None, or just rely on Python’s implicit return value, which is also None.

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 the following:

>>> import sys

>>> class StandardOutputRedirector:
...     def __init__(self, new_output):
...         self.new_output = new_output
...
...     def __enter__(self):
...         self.std_output = sys.stdout
...         sys.stdout = self.new_output
...
...     def __exit__(self, *_):
...         sys.stdout = self.std_output
...

This context manager takes a file object as an argument. 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 restore the standard output to its original value.

To use StandardOutputRedirector, you can do something like this:

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

The outer with statement provides the file object that you’ll 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!" to your screen. When you leave the inner with code block, the standard output is set back to its original value.

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

StandardOutputRedirector is a quick example of a context manager that doesn’t have a useful value to return from .__enter__(). In this example, you rely on Python’s implicit behavior of returning None.

Measuring Execution Time

Just like any 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 piece of code:

>>> from time import perf_counter, sleep

>>> class Timer:
...     def __enter__(self):
...         self.start = perf_counter()
...
...     def __exit__(self, *_):
...         end = perf_counter()
...         print(f"Elapsed time: {end - self.start:.4f} seconds")
...

>>> with Timer():
...     # The code to measure goes here...
...     sleep(0.5)
...
Elapsed time: 0.5051 seconds

In the .__enter__() method, you call time.perf_counter() to get the current timestamp at the beginning of a with code block and store it in the .start attribute. This value represents the context manager’s initial state.

Once the with block ends, .__exit__() gets called. The method gets the time at the end of the block and stores the value in a local variable end. Finally, it prints a message to inform you of the elapsed time.

Indenting Text

What if the resource you wanted to manage is the text indentation level in some kind of report generator application? For example, say that you want to write code that generates HTML output like the following:

<div>
  <p>
    Hello, World!
  </p>
</div>

This could be a great exercise to help you understand how context managers work. So, before you check out the implementation, you might take some time and try to produce this output with the help of a context manager that you can use as shown below:

with Indenter() as indenter:
    indenter.print("<div>")
    with indenter:
        indenter.print("<p>")
        with indenter:
            indenter.print("Hello, World!")
        indenter.print("</p>")
    indenter.print("</div>")

Notice how this code enters and leaves the same context manager multiple times to switch between different indentation levels.

Ready? Here’s how you can implement this functionality using a class-based context manager:

class Indenter:
    def __init__(self, width=2):
        self.indentation = " " * width
        self.level = -1

    def __enter__(self):
        self.level += 1
        return self

    def __exit__(self, *_):
        self.level -= 1

    def print(self, text):
        print(self.indentation * self.level + text)

Here, .__enter__() increments .level by 1 every time the flow of execution enters the context. The method also returns the current instance, self. In .__exit__(), you decrease .level, so the printed text steps back one indentation level each time you exit the context.

The key point here is that returning self from .__enter__() allows you to reuse the same context manager across several nested with statements. This changes the text indentation level every time you enter and exit a given context.

A good exercise for you at this point would be to write a function-based version of this context manager. Function-based context managers is what you’ll be learning about next.