Context managers and the with statement

In this lesson you will learn

Explain how the with statement calls __enter__ and __exit__
Write a context manager class that handles exceptions correctly
Create a generator-based manager with @contextmanager
Use ready-made managers such as suppress, redirect_stdout and ExitStack

A program opens a file, a database connection or a network socket, and then something goes wrong halfway — an exception, or an early return. If the resource is not released, files stay locked, connections pile up and data may never reach the disk. You met with open(...) in the lesson on files; now let's see what happens inside it and write our own context managers for any resource that must be set up and then reliably cleaned up.

What with really does

A context manager is an object with two methods. with manager as x: calls manager.__enter__() and binds its return value to x. Then the block runs. Whatever happens next — the block ends normally, or there is a return, a break or an exception — Python calls manager.__exit__(exc_type, exc, tb). A file object is its own context manager, and its __exit__ closes the file:

Python· Edit the code and run it again

with open('notes.txt', 'w', encoding='utf-8') as f:
    f.write('first line\n')
    print('inside:', f.closed)
print('after:', f.closed)

try:
    with open('notes.txt', encoding='utf-8') as f:
        text = f.read()
        raise ValueError('something went wrong')
except ValueError as e:
    print('error:', e)
print('closed even after the error:', f.closed)

▸ Expected output

inside: False
after: True
error: something went wrong
closed even after the error: True

Python turns a with statement into roughly the following code. Notice that __exit__ receives the exception and decides whether it should keep propagating:

Python

manager = open('notes.txt', encoding='utf-8')
f = manager.__enter__()
try:
    text = f.read()
except BaseException as e:
    if not manager.__exit__(type(e), e, e.__traceback__):
        raise
else:
    manager.__exit__(None, None, None)

A simplified version: the real mechanism also covers a few rare cases.

This is the same try/finally pattern you would otherwise have to write by hand in every place. A context manager packs the setup and cleanup logic into one reusable object, so the code that uses the resource stays short and simply cannot forget the cleanup.

Your own context manager class

The three arguments of __exit__ describe the exception: its type, the exception object and the traceback. If the block finished normally, all three are None. The return value matters: a true value suppresses the exception, while False or None lets it propagate. The manager below restores a dictionary if anything inside the block fails — like a tiny database transaction:

Python· Edit the code and run it again

class Transaction:
    def __init__(self, account):
        self.account = account

    def __enter__(self):
        self.backup = dict(self.account)
        return self.account

    def __exit__(self, exc_type, exc, tb):
        if exc_type is not None:
            self.account.clear()
            self.account.update(self.backup)
        return False

wallet = {'Aysel': 100, 'Murad': 50}
for amount in (30, 500):
    try:
        with Transaction(wallet) as w:
            w['Aysel'] -= amount
            w['Murad'] += amount
            if w['Aysel'] < 0:
                raise ValueError('not enough money')
    except ValueError as e:
        print('cancelled:', e)
    print(wallet)

▸ Expected output

{'Aysel': 70, 'Murad': 80}
cancelled: not enough money
{'Aysel': 70, 'Murad': 80}

The first transfer (30) went through. The second broke halfway: when the error was detected, 500 had already been taken from Aysel and added to Murad, but __exit__ restored the backup. The calling code still sees the exception, because __exit__ returned False.

contextlib.contextmanager: a manager from a generator

Writing a class for every manager is tedious. The decorator **@contextmanager** turns a generator function into a context manager: the code before yield plays the role of __enter__, the value after yield goes to as, and the code after yield plays the role of __exit__. If the block raises an exception, it is re-raised inside the generator right at the yield line — so the cleanup code must live in finally:

Python· Edit the code and run it again

from contextlib import contextmanager

@contextmanager
def tag(name):
    print(f'<{name}>')
    try:
        yield name
    finally:
        print(f'</{name}>')

with tag('ul'):
    for item in ['tea', 'plov']:
        with tag('li') as t:
            print(f'  {item} (inside {t})')

▸ Expected output

<ul>
<li>
  tea (inside li)
</li>
<li>
  plov (inside li)
</li>
</ul>

The output shows the order: each li opens and closes inside the ul, like nested brackets. Managers are exited in the reverse order of entering — last in, first out.

Example: a manager that times a block

Write a context manager timer(label) that prints how many milliseconds the code inside the with block took. The time must be printed even if the block raises an exception.

Show solution

Remember the start time before yield: start = time.perf_counter().
perf_counter() is a precise clock meant for measuring intervals.
Put yield inside try and compute time.perf_counter() - start in finally — then the time is printed even after an error.
The block does not need a value, so a bare yield is enough and as is not used.

Python· Edit the code and run it again

import time
from contextlib import contextmanager

@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f'{label}: {elapsed * 1000:.1f} ms')

with timer('sum of squares'):
    total = sum(n * n for n in range(1_000_000))
print(total)

The number of milliseconds depends on the computer, so it will differ a little every time.

Useful managers from the standard library

The contextlib module has several ready-made managers. suppress(Error) ignores the listed exceptions, and redirect_stdout(buffer) temporarily sends the output of print to another stream:

Python· Edit the code and run it again

from contextlib import suppress, redirect_stdout
import io
import os

with suppress(FileNotFoundError):
    os.remove('no-such-file.txt')
print('no crash')

buffer = io.StringIO()
with redirect_stdout(buffer):
    print('this goes into the buffer')
print('captured:', buffer.getvalue().strip())

▸ Expected output

no crash
captured: this goes into the buffer

When the number of resources is known only while the program runs, use **ExitStack**: it collects any number of managers and exits all of them, in reverse order, at the end of the block — even if opening the third file fails:

Python· Edit the code and run it again

from contextlib import ExitStack

names = ['a.txt', 'b.txt', 'c.txt']
for i, name in enumerate(names):
    with open(name, 'w', encoding='utf-8') as f:
        f.write(f'file {i}\n')

with ExitStack() as stack:
    files = [stack.enter_context(open(n, encoding='utf-8')) for n in names]
    print([f.readline().strip() for f in files])
print(all(f.closed for f in files))

▸ Expected output

['file 0', 'file 1', 'file 2']
True

Many other objects work with with too. Whenever something is set up temporarily, look for a context manager:

Object	What happens on exit
`open(...)`	the file is closed
`threading.Lock()`	the lock is released
`decimal.localcontext()`	the previous precision is restored
`tempfile.TemporaryDirectory()`	the folder and its contents are deleted
`unittest.mock.patch(...)`	the original object is put back
`sqlite3.connect(...)`	the transaction is committed or rolled back (the connection is not closed!)

Python· Edit the code and run it again

from decimal import Decimal, localcontext

print(Decimal(1) / Decimal(7))
with localcontext() as ctx:
    ctx.prec = 5
    print(Decimal(1) / Decimal(7))
print(Decimal(1) / Decimal(7))

▸ Expected output

0.1428571428571428571428571429
0.14286
0.1428571428571428571428571429

Inside the block the precision is 5 digits; on exit the default of 28 digits comes back by itself.

Exercise

Complete the context manager section(title): it prints == title == on entry and -- end of title -- on exit. The exit line must be printed even when the block raises an exception, and the exception must still reach the caller.

Exercise · Python· Edit the code and run it again

from contextlib import contextmanager

@contextmanager
def section(title):
    # print '== title ==' on entry and '-- end of title --' on exit, even after an error
    yield

with section('Report'):
    print('all good')

try:
    with section('Import'):
        raise ValueError('bad file')
except ValueError as e:
    print('error:', e)

▸ Expected output

== Report ==
all good
-- end of Report --
== Import ==
-- end of Import --
error: bad file

Exercise

Fix the __exit__ method of Ignore(*exceptions): if the exception in the block is one of the given types, suppress it; otherwise (and when there is no exception) return False.

Exercise · Python· Edit the code and run it again

class Ignore:
    def __init__(self, *exceptions):
        self.exceptions = exceptions

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        # return True only if the exception is one of self.exceptions
        return False

with Ignore(ZeroDivisionError):
    print(1 / 0)
print('still running')

with Ignore(KeyError, IndexError):
    [][5]
print('done')

▸ Expected output

still running
done

Key points

with m as x calls m.__enter__() (its result goes to x), and m.__exit__() is always called at the end.
__exit__(exc_type, exc, tb) receives the exception details; returning a true value suppresses the exception.
In @contextmanager, code before yield is the entry and code after it is the exit; cleanup belongs in finally.
Several managers are exited in reverse order; ExitStack handles a dynamic number of them.
suppress, redirect_stdout, localcontext and TemporaryDirectory are ready-made context managers.

Check yourself

10 questions. Every correct answer earns XP.

1 / 10

When is the __exit__ method called?