Using custom context in Python

If you’ve used python for a while you are familiar with the with statement, especially for reading and writing files. Using the with statement guarantees that the file will be closed even if an error occurs while executing the block of code inside the with statement This is kind of similar to using try... except... finally construct, except it's more cleaner and provides a more powerful abstraction. In Python, this is achieved through a context manager.

From python documentation

A context manager is an object that defines the runtime context to be established when executing a with statement. The context manager handles the entry into, and the exit from, the desired runtime context for the execution of the block of code. Context managers are normally invoked using the with statement (described in section The with statement), but can also be used by directly invoking their methods.

Typical uses of context managers include saving and restoring various kinds of global state, locking and unlocking resources, closing opened files, etc.

Defining a Custom context manager

To use the context manager behaviour in your object must define two required methods

1. __enter__(self)
2. __exit__(self, exc_type, exc_value, traceback)

__enter__(self) - You must return an object from this method. This function will be called by the with statement and the return value bound to the variable defined after as Usually you'll return self from this method.

__exit__(self, exc_type, exc_value, traceback) - This is where we do our clean up. For example closing connections, or releasing resource. This method will be called when the with block exits. If the code in the with block raises an exception, the exception information will be in the last 3 arguments - (exc_type, exc_value, traceback) of this function. If no exception occurred then these last 3 arguments will all be None. If an exception occurred in the with block, then you can either suppress the exception by returning a true value from this method. If you don't want to suppress errors then you you can return a value that evaluates to False.

Note: Default return value from a python function is None which evaluates to False. This means not returning anything from the __exit__ method will imply not supressing exceptions raised while executing the code in the with block.

Let's practice

Let's create a class for our custom context manage.

Create a new python file my_context.py

class MyContext:

    def __init__(self):
        # example file or database connection
        self.big_resource = "Some big resource"

    def __enter__(self):
       print('Starting context: ', self)
       return self

    def __exit__(self, exc_type, exc_value, traceback):
       print('Exiting context: ', self, exc_type, exc_value, traceback)
       # simulate cleaning up big_resource
       self.big_resource = None

       # suppress errors
       return True

    def do_something(self, data):
        # simulate doing something with big_resource
        print('Doing something with data: ', self, data)

Now let's use our context manager.

with MyContext() as f:
    data = 'Hello'
    f.do_something(data)

# assert resource is released
assert f.big_resource is None

Let's run our file from the terminal.

$ python my_context.py
Starting context <__main__.MyContext object at 0x10c21b5f8>
doing something with data Hello <__main__.MyContext object at 0x10c21b5f8>
Exiting context <__main__.MyContext object at 0x10c21b5f8> None None None

Let's make the with block raise excetion by calling the do_something function without an argument.

with MyContext() as f:
    data = 'Hello'
    f.do_something() # this should raise exception

assert f.big_resource is None

Let's run our file from the terminal.

$ python my_context.py
Starting context <__main__.MyContext object at 0x108e09f98>
Exiting context <__main__.MyContext object at 0x108e09f98> <class 'TypeError'> do_something() missing 1 required positional argument: 'data' <traceback object at 0x108e204c8>

Next time you are working with objects that need cleanup after using the object (such as database connections) try using a context manager. You can read more about about context managers in the python docs here

Happy coding!