Python Context Managers

摘自《Intermediate Python》

Sometimes, there is a need to execute some operations between another pair of operations. For example, open a file, read from the file and close the file or acquire a lock on a data structure, work with the data structure and release the data structure. These kinds of requirements come up most especially when dealing with system resources where the resource is acquired, worked with and then released. It is important that the acquisition and release of such resources are handled carefully so that any errors that may occur are correctly handled. Writing code to handle this all the time leads to a lot of repetition and cumbersome code. Context managers provide a solution to this. They provide a mean for abstracting away a pair of operations that are executed before and after another group of operation using the with statement.

The with statement enables a set of operations to run within a context. The context is controlled by a context manager object. An example of the use of the with statement is in opening files; this involves a pair of operations - opening and closing the file.

  # create a context 
    with open('output.txt', 'w') as f: 
        # carry out operations within context
        f.write('Hi there!')

The with statement can be used with any object that implements the context management protocol. This protocol defines a set of operations, __enter__ and __exit__ that are executed just before the start of execution of some piece of code and after the end of execution of some piece of code respectively. Generally, the definition and use of a context manager is shown in the following snippet.

  class context:
        def __enter__(self):
            set resource up
            return resource

        def __exit__(self, type, value, traceback):
            tear resource down

    # the context object returned by __enter__ method is bound to name
    with context() as name:
         do some functionality

If the initialised resource is used within the context then the__enter__ method must return the resource object so that it is bound within thewith statement using the as mechanism. A resource object must not be returned if the code being executed in the context doesn’t require a reference to the object that is set-up. The following is a very trivial example of a class that implements the context management protocol in a very simple fashion.

    >>> class Timer:
    ...     def __init__(self):
    ...         pass
    ...     def __enter__(self):
    ...         self.start_time = time.time()
    ...     def __exit__(self, type, value, traceback):
    ...         print("Operation took {} seconds to complete".format(time.time()-self.start_time))
    ... 

    ... 
    >>> with Foo():
    ...     print("Hey testing context managers")
    ... 
    Hey testing context managers
    Operation took 0.00010395050048828125 seconds to complete
    >>> 

When the with statement executes, the__enter__() method is called to create a new context; if a resource is initialized for use here then it is returned but this is not the case in this example. After the operations within the context are executed, the__exit__() method is called with the type, value and traceback as arguments. If no exception is raised during the execution of the of the operations within the context then all arguments are set to None. The__exit__ method returns a True or False depending on whether any raised exceptions have been handled. When False is returned then exception raised are propagated outside of the context for other code blocks to handle. Any resource clean-up is also carried out within the__exit__() method. This is all there is to context management. Now rather than write try...finally code to ensure that a file is closed or that a lock is released every time such resource is used, such chores can be handled in the the __exit__ method of a context manager class thus eliminating code duplication and making the code more intelligible.

The Contextlib module

For very simple use cases, there is no need to go through the hassle of implementing our own classes with __enter__ and __exit__ methods. The python contextlib module provides us with a high level method for implementing context manager. To define a context manager,the@contextmanager decorator from the contextlib module is used to decorate a function that handles the resource in question or carries out any initialization and clean-up; this function carrying out the initialization and tear down must however be a generator function. The following example illustrates this.

  from contextlib import contextmanager

    >>> from contextlib import contextmanager
    >>> @contextmanager
    ... def time_func():
    ...     start_time = time.time()
    ...     yield
    ...     print("Operation took {} seconds".format(time.time()-start_time))


    >>> with time_func():
    ...     print("Hey testing the context manager")
    ... 
    Hey testing the context manager
    Operation took 7.009506225585938e-05 seconds

This context generator function, time_func in this case, must yield exactly one value if it is required that a value be bound to a name in the with statement’s as clause. When generator yields, the code block nested in the with statement is executed. The generator is then resumed after the code block finishes execution. If an exception occurs during the execution of a block and is not handled in the block, the exception is re-raised inside the generator at the point where the yield occurred. If an exception is caught for purposes other than adequately handling such an exception then the generator must re-raise that exception otherwise the generator context manager will indicate to the with statement that the exception has been handled, and execution will resume normally after the context block.

Context managers just like decorators and metaclasses provide a clean method for abstracting away these kind of repetitive code that can clutter code and makes following code logic difficult.