Recursive Context Managers

read

Extending context managers to collections of dictionaries and lists:

What is a context manager?

A context manager is what allows you to use the with statement in Python. For example if you wanted to write to a file, you could do this:

f = open('test.txt', 'w')
f.write('hey!')
f.close()

assert f.closed is True  # this is True!

We always have to use open and close on files anyways, so context managers offer a convenient syntax using the with statement instead:

with open('test.txt', 'w') as f:
    f.write('hey')

assert f.closed is True  # this is also True!

The with statement will call f.close() for us automagically, so we just saved ourselves 1 line of code!

But how does it work?

Let’s implement a context manager to see how with works:

1
2
3
4
5
6
7
8
9
10
from contextlib import contextmanager

@contextmanager
def open_context_manager(name, mode):
    f = open(name, mode)
    yield f
    f.close()

with open_context_manager('test.txt', 'w') as f:
    f.write('hey!')

When you call with open_context_manager on line 9, the open_context_manager(...) function gets called on line 4 and gives you back the file handler f using the yield statement on line 6. Then it returns back to the code in the with block on line 10, and we write 'hey!' to the file with f.write('hey!'). Once we leave the with block, we return back to the open_context_manager(..) function on line 7, which calls f.close(). That’s it!

You could also implement a context manager with a try and finally statement:

def open_context_manager(name, mode):
    try:
        f = open(name, mode)
        yield f
    finally:
        f.close()

with open_context_manager('test.txt', 'w') as f:
    f.write('hey!')

Try it out! It works the same way.

So what was the problem?

The problem I had was that I needed to enter context managers all at the same time in a nested data structure.

If we wanted to enter a lot of context managers at the same time, we could start off by doing something like this:

with open('test1.txt', 'w') as f1,\
        open('test2.txt', 'w') as f2,\
        open('test3.txt', 'w') as f3:
    # do stuff here ...
    pass

But that isn’t really usable is it? Let’s say we had 100 files? I sure wouldn’t want to type that all out.

It turns out that all you need to implement a context manager in Python is to implement a class with an __enter__ and __exit__ method. __enter__ gets called when with is invoked on the object, and __exit__ gets called when you leave the with scope.

So here is how to implement a context manager for opening and closing files using a class:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Open:

    def __init__(self, name, mode):
        self.name = name
        self.mode = mode

    def __enter__(self):
        self.f = open(self.name, self.mode)
        return self.f

    def __exit__(self, *exc):
        self.f.close()

with Open('test.txt', 'w') as f:
    f.write('hey!')

It works the same way as before. __enter__ gets called when with Open(...) happens, and __exit__ gets called when you leave the with scope.

You might be wondering what the *exc is on line 11? It’s a list of required arguments for any __exit__ magic method, which get passed around so that you can handle exceptions gracefully. The arguments for *exc are exception_type, exception_value, and traceback.

So how can we leverage what we just did to enter multiple context managers at the same time? Let’s say we had files in a list, we could implement the following context manager:

1
2
3
4
5
6
7
8
9
10
11
12
class OpenFileList:

    def __init__(self, file_list, mode):
        self.file_list = file_list
        self.mode = mode

    def __enter__(self):
        self.fs = [open(f, self.mode) for f in self.file_list]
        return self.fs

    def __exit__(self, *exc):
        [f.close() for f in self.fs]

Boom! All we had to do was open files in a list on line 8, and then close them on line 12.

And now to use it:

with OpenFileList(['test.txt', 'test2.txt', 'test3.txt'], 'w') as fs:
    for f in fs:
        f.write('hey!')

with OpenFileList(['test.txt', 'test2.txt', 'test3.txt'], 'r') as fs:
    for f in fs:
        print(f.read())

Back to the original problem!

Unfortunately that doesn’t solve the original problem. I had multiple objects in a nested data structure made of lists and dicts. To solve the problem, we need to traverse the data strucutre and enter all nested contexts, then exit them later.

First we need a few useful functions:

Call a method on an object programmatically, so that we can do _apply_method(target, 'open', mode='r') instead of target.open('r'):

 from operator import methodcaller

 def _apply_method(f, method, *args, **kwargs):

     m = methodcaller(method, *args, **kwargs)

     if hasattr(f, method):
         return m(f)

     return f

We need to apply the method recursively so that we can enter/exit all nested context managers. Here we apply it recursively on nested dicts and lists:

 def _recursive_apply_method(f, method, *args, **kwargs):

     if isinstance(f, dict):
         return {k: _recursive_apply_method(v, method, *args, **kwargs)
                 for k, v in f.items()}
     elif isinstance(f, list):
         return [_recursive_apply_method(v, method, *args, **kwargs)
                 for v in f]

     return _apply_method(f, method, *args, **kwargs)

Using these two useful functions, we can now enter multiple nested contexts using a single context manager!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class recursive_file_context_manager(object):

    def __init__(self, files):
        self.files = files

    def __enter__(self):
        self.contexts = _recursive_apply_method(self.files, 'open', mode=self.mode)
        self.files = _recursive_apply_method(self.contexts, '__enter__')

        return self.files

    def __exit__(self, *exc):
        _recursive_apply_method(self.contexts, '__exit__', *exc)


class Read(recursive_file_context_manager):

    mode = 'r'


class Write(recursive_file_context_manager):

    mode = 'w'

We open all the file objects on line 7 by applying 'open' to each nested object. Then we enter the contexts we got back on line 8 by calling __enter__ on them. Finally, we exit the contexts on line 13 by calling __exit__ on all the nested contexts.

Sweet that wasn’t too bad. Now we could have some files in a nested structure like so:

from pathlib import Path

files = {
    'hey': Path('test1.txt'),
    'list': [Path('test2.txt'), Path('test3.txt')]
}

And we can enter all of them easily like this:

with Write(files) as fobj:
    fobj['hey'].write('yo!')
    fobj['list'][0].write('yo1!')
    fobj['list'][1].write('yo1!')

assert fobj['hey'].closed is True  # yay!

Recursive Context Managers

What is a context manager?

But how does it work?

So what was the problem?

Back to the original problem!

Written by

btaba

Supported by