Cleanup You Can Count On

Some things have to be undone: a file closed, a database connection closed, a temporary file deleted, a lock released. The hard part isn’t the undoing. It’s making sure it happens on every path out of the code: the normal end, an early return, a break, and above all an exception you didn’t expect.

That’s the job of with, and today you’ll write your own.

What with guarantees

You’ve been writing with open(path) as file: since the beginner course. Here’s what it actually promises, shown with a class that announces each step:

class Announcer:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"  enter {self.name}")
        return self                          # this is what "as" receives

    def __exit__(self, exc_type, exc, traceback):
        print(f"  exit  {self.name} (exception: {exc_type.__name__ if exc_type else None})")
        return False                         # False: don't swallow the exception


print("normal end:")
with Announcer("a"):
    print("  body")

print("return from inside:")
def early():
    with Announcer("b"):
        return "returned"
print(" ", early())

print("exception inside:")
try:
    with Announcer("c"):
        1 / 0
        print("  never printed")
except ZeroDivisionError:
    print("  the exception still reached the caller")

A with block calls __enter__ on the way in, and __exit__ on the way out, however it leaves. __exit__ is told whether an exception is on its way out, and returning False lets it carry on to the caller. Any object with these two methods is a context manager.

A trap: with conn doesn’t close

On Day 3 you used with conn: for transactions. It’s easy to assume it closes the connection, the way with open(...) closes a file. It doesn’t:

import sqlite3

with sqlite3.connect(":memory:") as conn:
    conn.execute("CREATE TABLE t (x)")
    conn.execute("INSERT INTO t VALUES (1)")

print(conn.execute("SELECT x FROM t").fetchall())    # still works: never closed
conn.close()

For a connection, with means “one transaction”: commit if the block succeeds, roll back if it raises. What a context manager does on exit is up to whoever wrote it, so it’s worth checking in the documentation rather than assuming. For a database file on disk, a connection left open keeps the file locked, and leaks one more open file every time the program forgets.

Writing one with contextlib

Writing a class with __enter__ and __exit__ works, but for most cases contextlib.contextmanager is shorter. You write a generator that yields exactly once:

from contextlib import contextmanager


@contextmanager
def open_store(path=":memory:"):
    """Connect, and close the connection when the with block ends, however it ends."""
    conn = connect(path)          # before yield: runs on the way in
    try:
        yield conn                # the value "as" receives; the with block runs here
    finally:
        conn.close()              # after yield: runs on the way out

with open_store("logs.db") as conn:
    save_records(conn, "server.log", records)

The with block runs where the generator yields. When the block ends, the generator carries on after the yield. If the block raised, the exception is raised inside the generator, at the yield, and that’s why the try and finally matter.

Bug hunt: the connection that stayed open

Here’s the first version of open_store, with a test that checks the connection is closed even when the block fails:

import sqlite3
import sys
import unittest
from contextlib import contextmanager


@contextmanager
def open_store(path=":memory:"):
    conn = sqlite3.connect(path)
    yield conn
    conn.close()


class TestOpenStore(unittest.TestCase):
    def test_closes_the_connection(self):
        with open_store() as conn:
            conn.execute("SELECT 1")
        with self.assertRaises(sqlite3.ProgrammingError):   # "Cannot operate on a closed database"
            conn.execute("SELECT 1")

    def test_closes_the_connection_even_when_the_block_fails(self):
        with self.assertRaises(ZeroDivisionError):
            with open_store() as conn:
                1 / 0
        with self.assertRaises(sqlite3.ProgrammingError):
            conn.execute("SELECT 1")


suite = unittest.defaultTestLoader.loadTestsFromTestCase(TestOpenStore)
unittest.TextTestRunner(stream=sys.stdout, verbosity=2).run(suite)

Reproduce. The normal case passes. The failing-block case fails with “ProgrammingError not raised”. The connection was still usable after the with, so it was never closed.

Inspect. Is the cleanup line running at all? Add a print just before it:

    yield conn
    print("closing")
    conn.close()

Run it. closing is printed once, for the first test. For the second, nothing.

Before you open this, make the second test pass without breaking the first. Then compare with ours, including why the line after yield didn’t run.

    conn = sqlite3.connect(path)
    try:
        yield conn
    finally:
        conn.close()

Your turn

import os
import sys
import unittest
from contextlib import contextmanager


@contextmanager
def working_file(path, text):
    """Create path containing text for the length of a with block, then delete it."""
    raise NotImplementedError("your code here")
    yield


class TestWorkingFile(unittest.TestCase):
    def test_the_file_exists_inside_the_block(self):
        with working_file("sample.log", "line one\n") as path:
            with open(path) as file:
                self.assertEqual(file.read(), "line one\n")

    def test_the_file_is_gone_afterwards(self):
        with working_file("sample.log", "x"):
            pass
        self.assertFalse(os.path.exists("sample.log"))

    def test_the_file_is_gone_even_when_the_block_fails(self):
        with self.assertRaises(ValueError):
            with working_file("sample.log", "x"):
                raise ValueError("something broke")
        self.assertFalse(os.path.exists("sample.log"))

    def test_an_existing_file_is_refused_and_left_alone(self):
        with open("keep.txt", "w") as file:
            file.write("precious")
        self.addCleanup(os.remove, "keep.txt")
        with self.assertRaises(FileExistsError):
            with working_file("keep.txt", "x"):
                pass
        with open("keep.txt") as file:
            self.assertEqual(file.read(), "precious")

    def test_a_failed_write_leaves_nothing_behind(self):
        with self.assertRaises(UnicodeEncodeError):          # a lone surrogate can't be UTF-8
            with working_file("sample.log", "\ud800"):
                pass
        self.assertFalse(os.path.exists("sample.log"))


suite = unittest.defaultTestLoader.loadTestsFromTestCase(TestWorkingFile)
unittest.TextTestRunner(stream=sys.stdout, verbosity=2).run(suite)

@contextmanager
def working_file(path, text):
    """Create path containing text for the length of a with block, then delete it."""
    file = open(path, "x", encoding="utf-8")     # "x": a new file only
    try:
        with file:
            file.write(text)
        yield path
    finally:
        os.remove(path)

What you can now do

• Say exactly what with guarantees, and check what a particular context manager does on exit, rather than assuming.
• Close a database connection deliberately, not just commit it.
• Write a context manager with @contextmanager, with its cleanup in finally.
• Test cleanup on the failure path, where it matters, not only the happy path.

Next: the Day 4 checkpoint, a configurable analyzer, and, if something really does repeat, a decorator.