Context Managers

Concepts covered: Context Managers

The gap between open and close

The file section called close() by hand. Between open() and close() sit the lines that do the work, and when 1 of them raises, Python jumps out of the block and close() never runs. The handle then stays open until Python gets around to discarding the file object, and nothing in the language promises when that will be. A job that opens a file per customer and keeps those objects alive, or runs on an interpreter that collects them late, leaks a handle on every error until the operating system refuses to open more. The finally block from the error section fixes this, and with is a shorter way to write the same guarantee.

f = open("plays.txt", "r")
try:
    text = f.read()
finally:
    f.close()
with open("plays.txt", "r") as f:
    text = f.read()

with calls open() and binds the result to f. After the indented block runs, it closes f, even if the block raised. The 2 programs do identical work; the second cannot forget the cleanup, because the cleanup is part of the statement.

import io

with io.StringIO("alpha\nbeta\n") as f:
    for line in f:
        print(line, end="")
    print("inside, closed:", f.closed)
print("after, closed:", f.closed)

Output:

alpha
beta
inside, closed: False
after, closed: True

f is still a name after the block, and the object it points to has been closed; f.closed reports the change the moment the block ends.

Objects that know how to clean up

with works on any object that defines 2 methods. __enter__ runs at the top of the block, and whatever it returns is what as binds. __exit__ runs when the block ends, receiving 3 arguments that describe the exception when 1 was raised and are all None otherwise. A file object defines both, and its __exit__ calls close(). A class of your own can define them for any resource that needs releasing.

class Timer:
    def __enter__(self):
        print("enter: start the clock")
        return self

    def __exit__(self, exc_type, exc, tb):
        print("exit: stop the clock")

with Timer() as t:
    print("working")
print("done")

Output:

enter: start the clock
working
exit: stop the clock
done

Cleanup when the block raises

The run below puts a division by zero inside the block and a try around the whole with. Watch the order: __exit__ prints before the handler outside does.

class Guard:
    def __enter__(self):
        print("enter")
        return self

    def __exit__(self, exc_type, exc, tb):
        print("exit sees", exc_type)

try:
    with Guard():
        print("working")
        1 / 0
except ZeroDivisionError:
    print("caught outside the with")
enter
working
exit sees <class 'ZeroDivisionError'>
caught outside the with

The division failed inside the block, so Python ran __exit__ first, passing the exception's class as exc_type, and only then let the ZeroDivisionError continue outward to the except wrapped around the with. Cleanup happened before the error was handled. An __exit__ that returns True tells Python the exception is dealt with, and nothing outside the with ever sees it; returning None, as this one does by having no return, lets the exception travel on.

Two resources in one with

A copy job needs a source and a destination open together. One with can hold both, separated by a comma, and closes them in reverse order when the block ends. An error on any line inside still closes both.

with open("plays.txt", "r") as src, open("plays.tmp", "w") as dst:
    for line in src:
        dst.write(line)

Practice

What does this program print before it ends?

class Res:
    def __enter__(self):
        print("open")
        return self

    def __exit__(self, exc_type, exc, tb):
        print("close")

with Res() as r:
    print("step 1")
    int("oops")
    print("step 2")
print("after")

with calls __exit__ with 4 arguments every time: the object as self, then the 3 exception slots, filled with None when the block finished normally. A header that lists only self accepts 1 of them, so the call fails after the indented block has already run, with TypeError: __exit__() takes 1 positional argument but 4 were given. The traceback follows the block's own prints, and the cleanup line inside __exit__ never prints.

class Lock:
    def __enter__(self):
        return self

    def __exit__(self):
        print("released")

About this lesson

This section is part of the Python Foundations: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.