Context Managers: The with Statement#

Why?#

  • Python is garbage collected

  • Actually, does not make any guarantees when resources are freed (though C-Python does refcounting, deterministically)

  • Usually not a problem with memory-only data (like list, dict, etc)

  • Want more deterministic behavior for other resources

  • Automatic cleanup ⟶ exception safety

  • ⟶ simplicity

Example: Open File#

  • The prototypical example

  • open() return value (a io.TextIOWrapper instance) can be used as a context manager

  • ⟶ with

with open('/etc/passwd') as f:
    for line in f:
        if 'jfasch' in line:
            print(line)

jfasch:x:1000:1000:Joerg Faschingbauer:/home/jfasch:/bin/bash

Without with, this would have to look more ugly:

try:
    f = open('/etc/passwd')
    for line in f:
        if 'jfasch' in line:
            print(line)
finally:
    f.close()                  # <--- gosh: what if open() failed?

jfasch:x:1000:1000:Joerg Faschingbauer:/home/jfasch:/bin/bash

Example: Temporary Directory#

  • Create a tar archive file in a temporary directory

  • Nested with blocks

  • ⟶ hard to get manual cleanup right

  • ⟶ with to the rescue

tar-tmpdir.py#
import tempfile
import shutil
import os
import tarfile

with tempfile.TemporaryDirectory() as tmpd:
    # create toplevel tar directory, and cram stuff in it
    subdir = tmpd + '/os-credentials'
    os.mkdir(subdir)
    shutil.copy('/etc/passwd', subdir)
    shutil.copy('/etc/group', subdir)

    # tar it
    tarname = tmpd + 'os-credentials.tar.bz2'
    with tarfile.open(tarname, 'w') as tf:
        tf.add(subdir, 'os-credentials')

    # copy tarfile into its final location
    shutil.copy(tarname, os.path.expandvars('$HOME/os-credentials.tar.bz2'))

Example: Multiple with Items#

  • with not contrained to only one managed object

  • Arbitrarily many objects possible

with open('/etc/passwd') as p, open('/etc/group') as g:
    # do something with p and g
    pass

Under The Hood: Context Manager#

  • Anything that has methods __enter__ and __exit__

  • __enter__: returns the target - the variable which is set by as

  • __exit__: cleans up resources, and receives exception context if any

    • Not called if __enter__ failed

    • Exception ignored if returns True

    • Exception re-raised if returns False (can omit return False ⟶ return None implicitly)

  • Example: manual open() context manager

  • (attention: complete nonsense because open() does that already)

class OpenForReading:
    def __init__(self, filename):
        self.filename = filename

    def __enter__(self):
        self.file = open(self.filename)
        return self.file      # <--- becomes 'f' in 'as f'
    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.file.close()
        return False          # <--- re-raise exception

with OpenForReading('/etc/passwd') as f:
    # do something with f
    raise RuntimeError('bad luck')

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[4], line 14
     10         return False          # <--- re-raise exception
     12 with OpenForReading('/etc/passwd') as f:
     13     # do something with f
---> 14     raise RuntimeError('bad luck')

RuntimeError: bad luck

Example: Monkeypatching The print Function#

class PrefixPrint:
    def __init__(self, prefix):
        self.prefix = prefix

    def myprint(self, *args, **kwargs):
        my_args = (self.prefix,) + args
        self.orig_print(*my_args, **kwargs)

    def __enter__(self):
        global print
        self.orig_print = print   # <--- save away original print
        print = self.myprint      # <--- override print

    def __exit__(self, exc_type, exc_val, exc_tb):
        global print
        print = self.orig_print   # <--- restore original print
        return False              # <--- re-raise exception if any

print('not cool')            # <--- prints: "not cool"
with PrefixPrint('MEGA:'):
    print('super cool')      # <--- prints: "MEGA: super cool"
print('not cool again')      # <--- prints: "not cool again"

not cool
MEGA: super cool
not cool again

Still Much Typing ⟶ @contextlib.contextmanager#

  • __enter__ and __exit__ still too clumsy

  • ⟶ using yield to split a function in half

  • Usually using try and finally for setup and teardown

  • Example: distilling OpenForReading() to a minimum

import contextlib

@contextlib.contextmanager
def OpenForReading(filename):
    file = open(filename)
    yield file            # <--- give control to with block ('file' becomes 'f' in 'as f')
    file.close()          # <--- continuing here after 'with' block has run

More Involved: Using Closures To Implement PrefixPrint#

import contextlib

@contextlib.contextmanager
def PrefixPrint(prefix):
    global print
    orig_print = print       # <--- save away original print

    def myprint(*args, **kwargs):
        myargs = (prefix,) + args
        orig_print(*myargs, **kwargs)

    print = myprint

    try:
        yield                # <--- give control to user's with block
    finally:
        print = orig_print   # <--- restore original print

print('not cool')            # <--- prints: "not cool"
with PrefixPrint('MEGA:'):
    print('super cool')      # <--- prints: "MEGA: super cool"
print('not cool again')      # <--- prints: "not cool again"

not cool
MEGA: super cool
not cool again