Defining the __enter__() and __exit__() methods

The defining feature of a context manager is that it has two special methods: __enter__() and __exit__(). These are used by the with statement to enter and exit the context. We'll use a simple context so that we can see how they work.

We'll often use context managers to make transient global changes. This might be a change to the database transaction status or a change to the locking status, something that we want to do and then undo when the transaction is complete.

For this example, we'll make a global change to the random number generator. We'll create a context in which the random number generator uses a fixed and known seed, providing a fixed sequence of values.

The following is the context manager class definition:

import random
class KnownSequence:
    def __init__(self, seed=0):
        self.seed= 0
    def __enter__(self):
        self.was= random.getstate()
        random.seed(self.seed, version=1)
        return self
    def __exit__(self, exc_type, exc_value, traceback):
        random.setstate(self.was)

We defined the required __enter__() and __exit__() methods. The __enter__() method will save the previous state of the random module and then reset the seed to a given value. The __exit__() method will restore the original state of the random number generator.

Note that __enter__() returns self. This is common for mixin context managers that have been added into other class definitions. We'll look at the concept of a mixin in Chapter 8, Decorators And Mixins – Cross-cutting Aspects.

The __exit__() method's parameters will have the value of None under normal circumstances. Unless we have specific exception-handling needs, we generally ignore the argument values. We'll look at exception-handling in the following code.

Here's an example of using the context:

print( tuple(random.randint(-1,36) for i in range(5)) )
with KnownSequence():
    print( tuple(random.randint(-1,36) for i in range(5)) )
print( tuple(random.randint(-1,36) for i in range(5)) )
with KnownSequence():
    print( tuple(random.randint(-1,36) for i in range(5)) )
print( tuple(random.randint(-1,36) for i in range(5)) )

Each time we create an instance of KnownSequence, we're modifying the way the random module works. During the context of the with statement, we'll get a fixed sequence of values. Outside the context, the random seed is restored, and we get random values.

The output will look like the following (in most cases):

(12, 0, 8, 21, 6)
(23, 25, 1, 15, 31)
(6, 36, 1, 34, 8)
(23, 25, 1, 15, 31)
(9, 7, 13, 22, 29)

Some of this output is machine-dependent. While the exact values may vary, the second and fourth lines will match because the seed was fixed by the context. The other lines will not necessarily match because they rely on the random module's own randomization features.