Search code examples
pythonpython-3.xgeneratorcontextmanager

Try finally block with generator in python


Can someone explain me the idea of generator and try except in this code:

from contextlib import contextmanager
 
@contextmanager
def file_open(path):
    try:
        f_obj = open(path, 'w')
        yield f_obj
    except OSError:
        print("We had an error!")
    finally:
        print('Closing file')
        f_obj.close()
 
 
if __name__ == '__main__':
    with file_open('test.txt') as fobj:
        fobj.write('Testing context managers')

As I know, finally is always executed regardless of correctness of the expression in try. So in my opinion this code should work like this: if we haven't exceptions, we open file, go to generator and the we go to finally block and return from the function. But I can't understand how generator works in this code. We used it only once and that's why we can't write all the text in the file. But I think my thoughts are incorrect. WHy?


Solution

  • So, one, your implementation is incorrect. You'll try to close the open file object even if it failed to open, which is a problem. What you need to do in this case is:

    @contextmanager
    def file_open(path):
        try:
            f_obj = open(path, 'w')
            try:
                yield f_obj
            finally:
                print('Closing file')
                f_obj.close()         
        except OSError:
            print("We had an error!")
    

    or more simply:

    @contextmanager
    def file_open(path):
        try:
            with open(path, 'w') as f_obj:
                yield f_obj
                print('Closing file')
        except OSError:
            print("We had an error!")
    

    To "how do generators in general work?" I'll refer you to the existing question on that topic. This specific case is complicated because using the @contextlib.contextmanager decorator repurposes generators for a largely unrelated purpose, using the fact that they innately pause in two cases:

    1. On creation (until the first value is requested)
    2. On each yield (when each subsequent value is requested)

    to implement context management.

    contextmanager just abuses this to make a class like this (actual source code is rather more complicated to cover edge cases):

    class contextmanager:
        def __init__(self, gen):
            self.gen = gen  # Receives generator in initial state
        def __enter__(self):
            return next(self.gen)  # Advances to first yield, returning the value it yields
        def __exit__(self, *args):
            if args[0] is not None:
                self.gen.throw(*args)  # Plus some complicated handling to ensure it did the right thing
            else:
                try:
                    next(self.gen)      # Check if it yielded more than once
                except StopIteration:
                    pass                # Expected to only yield once
                else:
                    raise RuntimeError(...)  # Oops, it yielded more than once, that's not supposed to happen
    

    allowing the coroutine elements of generators to back a simpler way to write simple context managers.