contextlib.py 4.04 KB
Newer Older
1 2 3
"""Utilities for with-statement contexts.  See PEP 343."""

import sys
4
from functools import wraps
5

6
__all__ = ["contextmanager", "nested", "closing"]
7

8 9
class GeneratorContextManager(object):
    """Helper for @contextmanager decorator."""
10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28

    def __init__(self, gen):
        self.gen = gen

    def __enter__(self):
        try:
            return self.gen.next()
        except StopIteration:
            raise RuntimeError("generator didn't yield")

    def __exit__(self, type, value, traceback):
        if type is None:
            try:
                self.gen.next()
            except StopIteration:
                return
            else:
                raise RuntimeError("generator didn't stop")
        else:
29 30 31 32
            if value is None:
                # Need to force instantiation so we can reliably
                # tell if we get the same exception back
                value = type()
33 34
            try:
                self.gen.throw(type, value, traceback)
35
                raise RuntimeError("generator didn't stop after throw()")
36
            except StopIteration, exc:
37 38 39
                # Suppress the exception *unless* it's the same exception that
                # was passed to throw().  This prevents a StopIteration
                # raised inside the "with" statement from being suppressed
40
                return exc is not value
41
            except:
42 43 44 45 46 47 48
                # only re-raise if it's *not* the exception that was
                # passed to throw(), because __exit__() must not raise
                # an exception unless __exit__() itself failed.  But throw()
                # has to raise the exception to signal propagation, so this
                # fixes the impedance mismatch between the throw() protocol
                # and the __exit__() protocol.
                #
49 50
                if sys.exc_info()[1] is not value:
                    raise
51 52


53 54
def contextmanager(func):
    """@contextmanager decorator.
55 56 57

    Typical usage:

58
        @contextmanager
59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
        def some_generator(<arguments>):
            <setup>
            try:
                yield <value>
            finally:
                <cleanup>

    This makes this:

        with some_generator(<arguments>) as <variable>:
            <body>

    equivalent to this:

        <setup>
        try:
            <variable> = <value>
            <body>
        finally:
            <cleanup>

    """
81
    @wraps(func)
82
    def helper(*args, **kwds):
83
        return GeneratorContextManager(func(*args, **kwds))
84 85 86
    return helper


87
@contextmanager
88
def nested(*managers):
89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105
    """Support multiple context managers in a single with-statement.

    Code like this:

        with nested(A, B, C) as (X, Y, Z):
            <body>

    is equivalent to this:

        with A as X:
            with B as Y:
                with C as Z:
                    <body>

    """
    exits = []
    vars = []
106
    exc = (None, None, None)
107
    try:
108 109 110 111 112 113 114 115
        for mgr in managers:
            exit = mgr.__exit__
            enter = mgr.__enter__
            vars.append(enter())
            exits.append(exit)
        yield vars
    except:
        exc = sys.exc_info()
116 117 118 119
    finally:
        while exits:
            exit = exits.pop()
            try:
120 121
                if exit(*exc):
                    exc = (None, None, None)
122 123 124
            except:
                exc = sys.exc_info()
        if exc != (None, None, None):
125 126 127 128
            # Don't rely on sys.exc_info() still containing
            # the right information. Another exception may
            # have been raised and caught by an exit method
            raise exc[0], exc[1], exc[2]
129 130


131 132
class closing(object):
    """Context to automatically close something at the end of a block.
133 134 135 136 137 138 139 140 141 142 143 144 145 146 147

    Code like this:

        with closing(<module>.open(<arguments>)) as f:
            <block>

    is equivalent to this:

        f = <module>.open(<arguments>)
        try:
            <block>
        finally:
            f.close()

    """
148 149 150 151 152 153
    def __init__(self, thing):
        self.thing = thing
    def __enter__(self):
        return self.thing
    def __exit__(self, *exc_info):
        self.thing.close()