How to Implement a Context Manager in Python


The primary purpose of a context manager is to manipulate how an object instance is controlled over a with statement.

The with statement was originally created to help facilitate the use of try/finally clauses. If you don't know what a try/finally clause is, check this link. In short, the try/finally clause assures us that code is executed.

The most common use of the with statement occurs in the opening of a file to read or write to it as seen below.

    try:

        with open('customers.json', 'r') as fl:
            decoded = json.load(fl)

    except ValueError as VE:
        print "Couldn't open the json emails", VE


In the example above, I'm opening a JSON file to load some contacts into my script. Starting at the with statement, the code is managed by the context manager. 

As I stated before, the context manager is designed to control the with statement. The context manager's operates, primarily, on two methods, the __enter__ and __exit__ methods. Both of these methods act as building and tearing-down methods, respectively.

To put the scenario in a very simple way, beginning at the with clause, the __enter__ method is invoked. After executing the code written after the with statement, a call to the __exit__ method occurs. The as clause doesn't always appear along the with clause, so you won't always need it. However, it is used with open to be able to reference the file.

>>> import json
>>> with open('contacts.json', 'r') as fil:
...     decoded = json.load(fil)
...
>>> fil
<closed file 'contacts.json', mode 'r' at 0x100496930>

Since the context manager governs the with statement, a context manager object is produced from the evaluation of the expressions after with. However, keep in mind that the value of the target variable (named 'fl' above) is not the context manager object, it's the result of calling the __enter__  method.
As you can see above, the context manager also takes care of closing the file by calling the __exit__ method at the end of the with block. 'fil' is now nothing but a 'closed file'.

Now that we have at least a vague image of how the context manager maintains secure code operation, Lets take a look under the curtains and see how a simple context manager is implemented with __enter__ and __exit__ methods.

class Numbers(object):

    def __init__(self):
        self.numbers = []

    def addNumber(self, number):
        self.numbers.append(number)

    def __enter__(self):
        return self.__dict__['numbers']

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is IndexError:
            print("Index Doesn't exist")
        elif exc_type is TypeError:
            print("Use proper Indexing")
        return True

The __enter__ method returns self.__dict__['numbers']. If you're not aware of what self.__dict__ is, its simply a dictionary that exhibits all the values initiated by the constructor. Each key in self.__dict__ corresponds to the variables created in __init__ (self.numbers above) and the values of each key are the values of the variables, which may change at any point based on your application. Take a look what I mean below.

>>> numbers = Numbers()
>>> numbers.addNumber(1)
>>> numbers.addNumber(2)
>>> numbers.__dict__
{'numbers': [1, 2]}

Now let's try invoking the __enter__ method by using the with statement. Keep in mind, as we said before, the __exit__ method is invoked automatically at the end of the with statement or if needed to before that, for example if an error is raised. In the __exit__ method above, I've included exception handling for IndexError and TypeError. The former handles an out of range index and the latter handles an incorrect index or an index that isn't an integer (str for example). Take a look at what I mean below.

>>> numbers = Numbers()
>>> numbers.addNumber(1)
>>> numbers.addNumber(2)
>>> numbers.__dict__
{'numbers': [1, 2]}
>>> with numbers as nm:
...     print nm[2]
...
Index Doesn't exist
>>> with numbers as nm:
...     print nm[1]
...
2
>>> with numbers as nm:
...     print nm['one']
...
Use proper Indexing

The only indexes available in self.__dict__ are 0 and 1 (at this particular moment of the instance). An attempt to extract a variable with index 2 will raise IndexError and an attempt to use a string to index will raise a TypeError.

Note: There's a workaround that would enable string indexing. Check out this post here

Moving along, Python's standard library also includes the contextlib module that allows you to build your own context manager as a function instead of a class. That's next.

@contextmanager

Python, through the contextlib module, supplies a built in generator function that lets you build your own context manager. This generator function comes in the form of a decorator. Luckily, the decorator makes the process of building the context manager much simpler constructing a whole class and implementing __enter__ and __exit__ methods.

As in most generator functions, the yield clause is utilized. However, since this particular generator function is wrapped with the @contextmanager decorator, the yield clause acts primarily to separate between the bodies of the __enter__ and __exit__ methods, which means it also acts to returns the output of the __enter__ method.

@contextlib.contextmanager
def Numbers(lst, index):
    try:
        yield lst[index]
    except (IndexError, TypeError):
        print 'Invalid Index'

The code block above is a replica of the class Numbers formulated within a function. It takes two arguments, 1) lst; a list of numbers and 2) index: the index you wish to return. The function is really simple yet it demonstrates the how the contextmanager decorator works. Take a look below at how we call this generator function.

>>> l = [1,2,4]
>>> index = 0
>>> with Numbers(l, index) as nums:
...     print nums
...
1
>>> index = 5
>>> with Numbers(l, index) as nums:
...     print nums
...
Wrong Index
>>> index = 'one'
>>> with Numbers(l, index) as nums:
...     print nums
...

Wrong Index