Python Context Manager

The purpose of python context manager objects is to control a with statement.
It is primarily used to make try/except/finally block simple.
The finally block exists to execute some code no matter what happens after try/except block, which is usually releasing some resources or clean-up.
Just like finally block, context manager objects actually do perform some block of code.

Context Manager Protocol

The context manager protocol consists of two methods: __enter__ and __exit__.
As you can imagine, __enter__ is called when the with statement starts, which is invoked on the context manager object.
The __exit__ method is called at the end of the with block.

Let’s take a look at a quick example.

class PrintName:
    """This class prints a string NYCOMDORI
    """

    def __enter__(self):
        print('NYCOMDORI ENTER')
        return 12345
    
    def __exit__(self, exc_type, exc_value, traceback):
        print('NYCOMDORI EXIT')
        if exc_type is not None:
            print("Exception happened")
            return True
            
>>> with PrintName() as pn:
	print(pn)

	
NYCOMDORI ENTER
12345
NYCOMDORI EXIT
>>> 

In the example, I just printed pn but you also see additional strings are printed.
I printed “ENTER” when __enter__ is called and “EXIT” when __exit__ is called.
Again, __enter__ is called automatically when with statement starts and __exit__ is called when with statement is done.

There is an important thing to note in this example.
The __enter__ method is invoked on the context manager object – PrintName.
However, pn variable is actually whatever __enter__ returned! In this case, it’s a number 12345.
It means that the context manager object is the result of evaluating the expression after with but pn is the result of calling __enter__ method.
It also implies that as clause is completely optional.

For the last, I just would like to briefly explain exception handling and the parameters in __exit__ method.
Please refer to the python document for more explanation.

please note that I explicitly returned True in order to suppress the exception.

The default behavior of handling exception for __exit__ method is:
1. check if a value is returned from __exit__ method.
2. if nothing is returned, essentially None will be returned, then it will assume exception is not handled properly and will propagate the exception.

The parameters are all related to handling exceptions within the with block.

exc_type: The exception class
exc_value: The exception instance
traceback: A traceback object

Here is another simple example that uses as clause – popular file open example.

>>> with open('test.txt') as f:
	f.read(20)

'This is a test file'
>>> f
<_io.TextIOWrapper name='test.txt' mode='r' encoding='UTF-8'>
>>> f.encoding
'UTF-8'
>>> f.closed
True
>>> f.read(20)
Traceback (most recent call last):
  File "<pyshell#9>", line 1, in <module>
    f.read(20)
ValueError: I/O operation on closed file.

As you can see, although f is still available after with block you cannot do any I/O operation since it’s closed.
And now we all know that f is the result of calling open().

@contextmanager decorator

Although context manager objects could be very useful and handy you might feel lazy to implement __enter__ and __exit__ methods every time.
To make your life simpler, contextlib utilities actually provides a convenient decorator for you, which is called @contextmanager.

The decorator actually lets you build a context manager object from a generator function that you implement.
I understand this might be a little confusing at first glance because generator functions are usually used for iterators or generator expressions.
However, in this case, @contextmanager decorator only requires a single yield because it is only used to split the body of the functions in two parts.
All the expressions before the yield will be used when __enter__ method in the decorator is called.
And the other half, after the yield, will be called when __exit__ method is called.

It is actually quite convenient util provided by the python standard library.
It will be very obvious once you look at the example.

from contextlib import contextmanager

@contextmanager
def print_name():
    print('NYCOMDORI ENTER')
    try:
        yield 12345
    except:
        print('exception happened')
    print('NYCOMDORI EXIT')
    
>>> with print_name() as pn:
	print(pn)

	
NYCOMDORI ENTER
12345
NYCOMDORI EXIT

Using @contextmanager decorator yielded exactly the same result as PrintName class with a much simpler implementation.
Let me add some more explanation about the above example.

The generator function you implemented will be wrapped by the decorator.
The wrapper class actually implemented __enter__ and __exit__ method which calls part of the generator function when the methods are invoked.
As __enter__ method is invoked, it will call the generator function and holds the result to a variable (value returned by yield statement).
Then, it will call next function to invoke yield of which value will be bounded to a target variable.
Then, it will execute whatever is in with block.

The __exit__ method will be called once with block is done, which will first check if an exception was passed (as exc_type).
If that’s true, it will raise the exception in the yield line inside the generator function.
If an exception was not passed in, call next to resume executing the rest of the code after yield in the generator function.

That’s why I guarded the yield line with try/except to protect from an exception raised inside with block.
Please note that nothing is returned when handling exceptions.
It’s because by default @contextmanager assumes the exception is handled and should be suppressed.
You need to explicitly re-raise the exception if necessary.

Conclusion

In this post, I explained about context manager objects and how to implement them in two ways – one explicitly implementing the class with __enter__ and __exit__ methods and the other using @contextmanager decorator which conveniently implements the methods for you.
With context manager objects demystified you should feel more comfortable with using it.
Thank you for reading the post!