Context Managers and Python's with Statement

Context Managers and Python's with Statement

by Leodanis Pozo Ramos Jun 02, 2021 intermediate python

The with statement in Python is a quite useful tool for properly managing external resources in your programs. It allows you to take advantage of existing context managers to automatically handle the setup and teardown phases whenever you’re dealing with external resources or with operations that require those phases.

Besides, the context management protocol allows you to create your own context managers so you can customize the way you deal with system resources. So, what’s the with statement good for?

In this tutorial, you’ll learn:

  • What the Python with statement is for and how to use it
  • What the context management protocol is
  • How to implement your own context managers

With this knowledge, you’ll write more expressive code and avoid resource leaks in your programs. The with statement helps you implement some common resource management patterns by abstracting their functionality and allowing them to be factored out and reused.

Managing Resources in Python

One common problem you’ll face in programming is how to properly manage external resources, such as files, locks, and network connections. Sometimes, a program will retain those resources forever, even if you no longer need them. This kind of issue is called a memory leak because the available memory gets reduced every time you create and open a new instance of a given resource without closing an existing one.

Managing resources properly is often a tricky problem. It requires both a setup phase and a teardown phase. The latter phase requires you to perform some cleanup actions, such as closing a file, releasing a lock, or closing a network connection. If you forget to perform these cleanup actions, then your application keeps the resource alive. This might compromise valuable system resources, such as memory and network bandwidth.

For example, a common problem that can arise when developers are working with databases is when a program keeps creating new connections without releasing or reusing them. In that case, the database back end can stop accepting new connections. This might require an admin to log in and manually kill those stale connections to make the database usable again.

Another frequent issue shows up when developers are working with files. Writing text to files is usually a buffered operation. This means that calling .write() on a file won’t immediately result in writing text to the physical file but to a temporary buffer. Sometimes, when the buffer isn’t full and developers forget to call .close(), part of the data can be lost forever.

Another possibility is that your application runs into errors or exceptions that cause the control flow to bypass the code responsible for releasing the resource at hand. Here’s an example in which you use open() to write some text to a file:

file = open("hello.txt", "w")
file.write("Hello, World!")
file.close()

This implementation doesn’t guarantee the file will be closed if an exception occurs during the .write() call. In this case, the code will never call .close(), and therefore your program might leak a file descriptor.

In Python, you can use two general approaches to deal with resource management. You can wrap your code in:

  1. A tryfinally construct
  2. A with construct

The first approach is quite general and allows you to provide setup and teardown code to manage any kind of resource. However, it’s a little bit verbose. Also, what if you forget any cleanup actions?

The second approach provides a straightforward way to provide and reuse setup and teardown code. In this case, you’ll have the limitation that the with statement only works with context managers. In the next two sections, you’ll learn how to use both approaches in your code.

The tryfinally Approach

Working with files is probably the most common example of resource management in programming. In Python, you can use a tryfinally 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 tryexcept 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 tryfinally statement.

Compared to traditional tryfinally 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 tryfinally 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 tryfinally 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 tryfinally statement with setup and teardown code each time, with takes care of that for you and avoids repetition.

Using the Python with Statement

As long as Python developers have incorporated the with statement into their coding practice, the tool has been shown to have several valuable use cases. More and more objects in the Python standard library now provide support for the context management protocol so you can use them in a with statement.

In this section, you’ll code some examples that show how to use the with statement with several classes both in the standard library and in third-party libraries.

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 tryexcept 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.

Summarizing the with Statement’s Advantages

To summarize what you’ve learned so far, here’s an inexhaustive list of the general benefits of using the Python with statement in your code:

  • Makes resource management safer than its equivalent tryfinally statements
  • Encapsulates standard uses of tryfinally statements in context managers
  • Allows reusing the code that automatically manages the setup and teardown phases of a given operation
  • Helps avoid resource leaks

Using the with statement consistently can improve the general quality of your code and make it safer by preventing resource leak problems.

Using the async with Statement

The with statement also has an asynchronous version, async with. You can use it to write context managers that depend on asynchronous code. It’s quite common to see async with in that kind of code, as many IO operations involve setup and teardown phases.

For example, say you need to code an asynchronous function to check if a given site is online. To do that, you can use aiohttp, asyncio, and async with like this:

 1# site_checker_v0.py
 2
 3import aiohttp
 4import asyncio
 5
 6async def check(url):
 7    async with aiohttp.ClientSession() as session:
 8        async with session.get(url) as response:
 9            print(f"{url}: status -> {response.status}")
10            html = await response.text()
11            print(f"{url}: type -> {html[:17].strip()}")
12
13async def main():
14    await asyncio.gather(
15        check("https://realpython.com"),
16        check("https://pycoders.com"),
17    )
18
19asyncio.run(main())

Here’s what this script does:

  • Line 3 imports aiohttp, which provides an asynchronous HTTP client and server for asyncio and Python. Note that aiohttp is a third-party package that you can install by running python -m pip install aiohttp on your command line.
  • Line 4 imports asyncio, which allows you to write concurrent code using the async and await syntax.
  • Line 6 defines check() as an asynchronous function using the async keyword.

Inside check(), you define two nested async with statements:

  • Line 7 defines an outer async with that instantiates aiohttp.ClientSession() to get a context manager. It stores the returned object in session.
  • Line 8 defines an inner async with statement that calls .get() on session using url as an argument. This creates a second context manager and returns a response.
  • Line 9 prints the response status code for the url at hand.
  • Line 10 runs an awaitable call to .text() on response and stores the result in html.
  • Line 11 prints the site url and its document type, doctype.
  • Line 13 defines the script’s main() function, which is also a coroutine.
  • Line 14 calls gather() from asyncio. This function runs awaitable objects in a sequence concurrently. In this example, gather() runs two instances of check() with a different URL for each.
  • Line 19 runs main() using asyncio.run(). This function creates a new asyncio event loop and closes it at the end of the operation.

If you run this script from your command line, then you get an output similar to the following:

$ python site_checker_v0.py
https://realpython.com: status -> 200
https://pycoders.com: status -> 200
https://pycoders.com: type -> <!doctype html>
https://realpython.com: type -> <!doctype html>

Cool! Your script works and you confirm that both sites are currently available. You also retrieve the information regarding document type from each site’s home page.

The async with statement works similar to the regular with statement, but it requires an asynchronous context manager. In other words, it needs a context manager that is able to suspend execution in its enter and exit methods. Asynchronous context managers implement the special methods .__aenter__() and .__aexit__(), which correspond to .__enter__() and .__exit__() in a regular context manager.

The async with ctx_mgr construct implicitly uses await ctx_mgr.__aenter__() when entering the context and await ctx_mgr.__aexit__() when exiting it. This achieves async context manager behavior seamlessly.

Creating Custom Context Managers

You’ve already worked with context managers from the standard library and third-party libraries. There’s nothing special or magical about open(), threading.Lock, decimal.localcontext(), or the others. They just return objects that implement the context management protocol.

You can provide the same functionality by implementing both the .__enter__() and the .__exit__() special methods in your class-based context managers. You can also create custom function-based context managers using the contextlib.contextmanager decorator from the standard library and an appropriately coded generator function.

In general, context managers and the with statement aren’t limited to resource management. They allow you to provide and reuse common setup and teardown code. In other words, with context managers, you can perform any pair of operations that needs to be done before and after another operation or procedure, such as:

  • Open and close
  • Lock and release
  • Change and reset
  • Create and delete
  • Enter and exit
  • Start and stop
  • Setup and teardown

You can provide code to safely manage any of these pairs of operations in a context manager. Then you can reuse that context manager in with statements throughout your code. This prevents errors and reduces repetitive boilerplate code. It also makes your APIs safer, cleaner, and more user-friendly.

In the next two sections, you’ll learn the basics of creating class-based and function-based context managers.

Coding Class-Based Context Managers

To implement the context management protocol and create class-based context managers, you need to add both the .__enter__() and the __exit__() special methods to your classes. The table below summarizes how these methods work, the arguments they take, and the logic you can put in them:

Method Description
.__enter__(self) This method handles the setup logic and is called when entering a new with context. Its return value is bound to the with target variable.
.__exit__(self, exc_type, exc_value, exc_tb) This method handles the teardown logic and is called when the flow of execution leaves the with context. If an exception occurs, then exc_type, exc_value, and exc_tb hold the exception type, value, and traceback information, respectively.

When the with statement executes, it calls .__enter__() on the context manager object to signal that you’re entering into a 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. If no exception occurs in the with code block, then the three last arguments to .__exit__() are set to None. Otherwise, they hold the type, value, and traceback associated with the exception at hand.

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 this feature to encapsulate exception handling inside the context manager.

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.

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.

Creating Function-Based Context Managers

Python’s generator functions and the contextlib.contextmanager decorator provide an alternative and convenient way to implement the context management protocol. If you decorate an appropriately coded generator function with @contextmanager, then you get a function-based context manager that automatically provides both required methods, .__enter__() and .__exit__(). This can make your life more pleasant by saving you some boilerplate code.

The general pattern to create a context manager using @contextmanager along with a generator function goes like this:

>>>
>>> from contextlib import contextmanager

>>> @contextmanager
... def hello_context_manager():
...     print("Entering the context...")
...     yield "Hello, World!"
...     print("Leaving the context...")
...

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

In this example, you can identify two visible sections in hello_context_manager(). Before the yield statement, you have the setup section. There, you can place the code that acquires the managed resources. Everything before the yield runs when the flow of execution enters the context.

After the yield statement, you have the teardown section, in which you can release the resources and do the cleanup. The code after yield runs at the end of the with block. The yield statement itself provides the object that will be assigned to the with target variable.

This implementation and the one that uses the context management protocol are practically equivalent. Depending on which one you find more readable, you might prefer one over the other. A downside of the function-based implementation is that it requires an understanding of advanced Python topics, such as decorators and generators.

The @contextmanager decorator reduces the boilerplate required to create a context manager. Instead of writing a whole class with .__enter__() and .__exit__() methods, you just need to implement a generator function with a single yield that produces whatever you want .__enter__() to return.

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.

Writing Good APIs With Context Managers

Context managers are quite flexible, and if you use the with statement creatively, then you can define convenient APIs for your classes, modules, and packages.

For example, what if the resource you wanted to manage is the text indentation level in some kind of report generator application? In that case, you could write code like this:

with Indenter() as indent:
    indent.print("hi!")
    with indent:
        indent.print("hello")
        with indent:
            indent.print("bonjour")
    indent.print("hey")

This almost reads like a domain-specific language (DSL) for indenting text. Also, notice how this code enters and leaves the same context manager multiple times to switch between different indentation levels. Running this code snippet leads to the following output and prints neatly formatted text:

hi!
    hello
        bonjour
hey

How would you implement a context manager to support this functionality? This could be a great exercise to wrap your head around how context managers work. So, before you check out the implementation below, you might take some time and try to solve this by yourself as a learning exercise.

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

class Indenter:
    def __init__(self):
        self.level = -1

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

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.level -= 1

    def print(self, text):
        print("    " * 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 every time you exit the context.

The key point in this example 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 leave a given context.

A good exercise for you at this point would be to write a function-based version of this context manager. Go ahead and give it a try!

Creating an Asynchronous Context Manager

To create an asynchronous context manager, you need to define the .__aenter__() and .__aexit__() methods. The script below is a reimplementation of the original script site_checker_v0.py you saw before, but this time you provide a custom asynchronous context manager to wrap the session creation and closing functionalities:

# site_checker_v1.py

import aiohttp
import asyncio

class AsyncSession:
    def __init__(self, url):
        self._url = url

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        response = await self.session.get(self._url)
        return response

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        await self.session.close()

async def check(url):
    async with AsyncSession(url) as response:
        print(f"{url}: status -> {response.status}")
        html = await response.text()
        print(f"{url}: type -> {html[:17].strip()}")

async def main():
    await asyncio.gather(
        check("https://realpython.com"),
        check("https://pycoders.com"),
    )

asyncio.run(main())

This script works similar to its previous version, site_checker_v0.py. The main difference is that, in this example, you extract the logic of the original outer async with statement and encapsulate it in AsyncSession.

In .__aenter__(), you create an aiohttp.ClientSession(), await the .get() response, and finally return the response itself. In .__aexit__(), you close the session, which corresponds to the teardown logic in this specific case. Note that .__aenter__() and .__aexit__() must return awaitable objects. In other words, you must define them with async def, which returns a coroutine object that is awaitable by definition.

If you run the script from your command line, then you get an output similar to this:

$ python site_checker_v1.py
https://realpython.com: status -> 200
https://pycoders.com: status -> 200
https://realpython.com: type -> <!doctype html>
https://pycoders.com: type -> <!doctype html>

Great! Your script works just like its first version. It sends GET requests to both sites concurrently and processes the corresponding responses.

Finally, a common practice when you’re writing asynchronous context managers is to implement the four special methods:

  1. .__aenter__()
  2. .__aexit__()
  3. .__enter__()
  4. .__exit__()

This makes your context manager usable with both variations of with.

Conclusion

The Python with statement is a powerful tool when it comes to managing external resources in your programs. Its use cases, however, aren’t limited to resource management. You can use the with statement along with existing and custom context managers to handle the setup and teardown phases of a given process or operation.

The underlying context management protocol allows you to create custom context managers and factor out the setup and teardown logic so you can reuse them in your code.

In this tutorial, you learned:

  • What the Python with statement is for and how to use it
  • What the context management protocol is
  • How to implement your own context managers

With this knowledge, you’ll write safe, concise, and expressive code. You’ll also avoid resource leaks in your programs.

🐍 Python Tricks 💌

Get a short & sweet Python Trick delivered to your inbox every couple of days. No spam ever. Unsubscribe any time. Curated by the Real Python team.

Python Tricks Dictionary Merge

About Leodanis Pozo Ramos

Leodanis Pozo Ramos Leodanis Pozo Ramos

Leodanis is an industrial engineer who loves Python and software development. He's a self-taught Python developer with 6+ years of experience. He's an avid technical writer with a growing number of articles published on Real Python and other sites.

» More about Leodanis

Each tutorial at Real Python is created by a team of developers so that it meets our high quality standards. The team members who worked on this tutorial are:

Master Real-World Python Skills With Unlimited Access to Real Python

Join us and get access to hundreds of tutorials, hands-on video courses, and a community of expert Pythonistas:

Level Up Your Python Skills »

Master Real-World Python Skills
With Unlimited Access to Real Python

Join us and get access to hundreds of tutorials, hands-on video courses, and a community of expert Pythonistas:

Level Up Your Python Skills »

What Do You Think?

Real Python Comment Policy: The most useful comments are those written with the goal of learning from or helping out other readers—after reading the whole article and all the earlier comments. Complaints and insults generally won’t make the cut here.

What’s your #1 takeaway or favorite thing you learned? How are you going to put your newfound skills to use? Leave a comment below and let us know.

Keep Learning

Related Tutorial Categories: intermediate python