Lib/contextlib.py

   1 """Utilities for with-statement contexts.  See PEP 343."""
   2
   3 import sys
   4
   5 __all__ = ["contextmanager", "nested", "closing"]
   6
   7 class GeneratorContextManager(object):
   8     """Helper for @contextmanager decorator."""
   9
  10     def __init__(self, gen):
  11         self.gen = gen
  12
  13     def __context__(self):
  14         return self
  15
  16     def __enter__(self):
  17         try:
  18             return self.gen.next()
  19         except StopIteration:
  20             raise RuntimeError("generator didn't yield")
  21
  22     def __exit__(self, type, value, traceback):
  23         if type is None:
  24             try:
  25                 self.gen.next()
  26             except StopIteration:
  27                 return
  28             else:
  29                 raise RuntimeError("generator didn't stop")
  30         else:
  31             try:
  32                 self.gen.throw(type, value, traceback)
  33                 return True
  34             except StopIteration:
  35                 return True
  36
  37
  38 def contextmanager(func):
  39     """@contextmanager decorator.
  40
  41     Typical usage:
  42
  43         @contextmanager
  44         def some_generator(<arguments>):
  45             <setup>
  46             try:
  47                 yield <value>
  48             finally:
  49                 <cleanup>
  50
  51     This makes this:
  52
  53         with some_generator(<arguments>) as <variable>:
  54             <body>
  55
  56     equivalent to this:
  57
  58         <setup>
  59         try:
  60             <variable> = <value>
  61             <body>
  62         finally:
  63             <cleanup>
  64
  65     """
  66     def helper(*args, **kwds):
  67         return GeneratorContextManager(func(*args, **kwds))
  68     try:
  69         helper.__name__ = func.__name__
  70         helper.__doc__ = func.__doc__
  71     except:
  72         pass
  73     return helper
  74
  75
  76 @contextmanager
  77 def nested(*contexts):
  78     """Support multiple context managers in a single with-statement.
  79
  80     Code like this:
  81
  82         with nested(A, B, C) as (X, Y, Z):
  83             <body>
  84
  85     is equivalent to this:
  86
  87         with A as X:
  88             with B as Y:
  89                 with C as Z:
  90                     <body>
  91
  92     """
  93     exits = []
  94     vars = []
  95     exc = (None, None, None)
  96     try:
  97         try:
  98             for context in contexts:
  99                 mgr = context.__context__()
 100                 exit = mgr.__exit__
 101                 enter = mgr.__enter__
 102                 vars.append(enter())
 103                 exits.append(exit)
 104             yield vars
 105         except:
 106             exc = sys.exc_info()
 107     finally:
 108         while exits:
 109             exit = exits.pop()
 110             try:
 111                 if exit(*exc):
 112                     exc = (None, None, None)
 113             except:
 114                 exc = sys.exc_info()
 115         if exc != (None, None, None):
 116             raise
 117
 118
 119 @contextmanager
 120 def closing(thing):
 121     """Context manager to automatically close something at the end of a block.
 122
 123     Code like this:
 124
 125         with closing(<module>.open(<arguments>)) as f:
 126             <block>
 127
 128     is equivalent to this:
 129
 130         f = <module>.open(<arguments>)
 131         try:
 132             <block>
 133         finally:
 134             f.close()
 135
 136     """
 137     try:
 138         yield thing
 139     finally:
 140         thing.close()