Propagate nested Python contexts

Recently I've been playing around with Python's context managers. Context managers are used in combination with the 'with' statement. As an example a file can be opened and closed in a context.
The file will be open in the 'with' block and closed  after the 'with' block.

with file('filename.txt') as f:
    for line in f:
        print(f)

To implement a context manager you'll have to create a class that inherits from object with a '__enter__' and '__exit__' method:

class Context(object):
    def __enter__(self):
        """Enter context"""

    def __exit__(self):
        """Exit context"""

Note that the syntax 'with <context_manager> [as <target>]:' assigns a target variable. This is the returned value of the '__enter__' method.

In my case I wanted to be able to nest contexts and ensure that nested contexts were propagated to the root context. My first idea was to have a class attribute and assign that in the '__enter__' method and return that value (even in nested contexts). This was the result:

class Context(object):
    current = None

    def __enter__(self):
        """Enter context"""
        if not Context.current:
            Context.current = self
        return Context.current

    def __exit__(self):
        """Exit context"""
        Context.current = None

After a quick test I noticed a nested context set the current context as None upon exit (the third 'print' statement will print None):

with Context() as root:
    print(Context.current)
    with Context() as nested:
        print(Context.current)
    print(Context.current)

To prevent this I modified the '__exit__' method to check if the '__exit__' is called for the root context. The final result looks like this:

class Context(object):
    current = None

    def __enter__(self):
        """Enter context"""
        if not Context.current:
            Context.current = self
        return Context.current

    def __exit__(self):
        """Exit context"""
        if self is Context.current:
            Context.current = None

By using the 'is' comparison the context only gets closed when the '__exit__' method is called for the root context manager.