blob: e90fc4181a39b406a94abddc409eede03fd8fd67 [file] [log] [blame]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00001"""Utilities for with-statement contexts. See PEP 343."""
2
3import sys
Christian Heimes81ee3ef2008-05-04 22:42:01 +00004from functools import wraps
Raymond Hettinger91e3b9d2009-05-28 22:20:03 +00005from warnings import warn
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00006
Michael Foordb3a89842010-06-30 12:17:50 +00007__all__ = ["contextmanager", "closing", "ContextDecorator"]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00008
Michael Foordb3a89842010-06-30 12:17:50 +00009
10class ContextDecorator(object):
11 "A base class or mixin that enables context managers to work as decorators."
Nick Coghlan0ded3e32011-05-05 23:49:25 +100012
13 def _recreate_cm(self):
14 """Return a recreated instance of self.
15
16 Allows otherwise one-shot context managers like
17 _GeneratorContextManager to support use as
18 decorators via implicit recreation.
19
20 Note: this is a private interface just for _GCM in 3.2 but will be
21 renamed and documented for third party use in 3.3
22 """
23 return self
24
Michael Foordb3a89842010-06-30 12:17:50 +000025 def __call__(self, func):
26 @wraps(func)
27 def inner(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +100028 with self._recreate_cm():
Michael Foordb3a89842010-06-30 12:17:50 +000029 return func(*args, **kwds)
30 return inner
31
32
Antoine Pitrou67b212e2011-01-08 09:55:31 +000033class _GeneratorContextManager(ContextDecorator):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000034 """Helper for @contextmanager decorator."""
35
Nick Coghlan0ded3e32011-05-05 23:49:25 +100036 def __init__(self, func, *args, **kwds):
37 self.gen = func(*args, **kwds)
38 self.func, self.args, self.kwds = func, args, kwds
39
40 def _recreate_cm(self):
41 # _GCM instances are one-shot context managers, so the
42 # CM must be recreated each time a decorated function is
43 # called
44 return self.__class__(self.func, *self.args, **self.kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000045
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000046 def __enter__(self):
47 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000048 return next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000049 except StopIteration:
50 raise RuntimeError("generator didn't yield")
51
52 def __exit__(self, type, value, traceback):
53 if type is None:
54 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000055 next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000056 except StopIteration:
57 return
58 else:
59 raise RuntimeError("generator didn't stop")
60 else:
Guido van Rossum2cc30da2007-11-02 23:46:40 +000061 if value is None:
62 # Need to force instantiation so we can reliably
63 # tell if we get the same exception back
64 value = type()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000065 try:
66 self.gen.throw(type, value, traceback)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000067 raise RuntimeError("generator didn't stop after throw()")
Guido van Rossumb940e112007-01-10 16:19:56 +000068 except StopIteration as exc:
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000069 # Suppress the exception *unless* it's the same exception that
70 # was passed to throw(). This prevents a StopIteration
71 # raised inside the "with" statement from being suppressed
72 return exc is not value
73 except:
74 # only re-raise if it's *not* the exception that was
75 # passed to throw(), because __exit__() must not raise
76 # an exception unless __exit__() itself failed. But throw()
77 # has to raise the exception to signal propagation, so this
78 # fixes the impedance mismatch between the throw() protocol
79 # and the __exit__() protocol.
80 #
81 if sys.exc_info()[1] is not value:
82 raise
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000083
84
85def contextmanager(func):
86 """@contextmanager decorator.
87
88 Typical usage:
89
90 @contextmanager
91 def some_generator(<arguments>):
92 <setup>
93 try:
94 yield <value>
95 finally:
96 <cleanup>
97
98 This makes this:
99
100 with some_generator(<arguments>) as <variable>:
101 <body>
102
103 equivalent to this:
104
105 <setup>
106 try:
107 <variable> = <value>
108 <body>
109 finally:
110 <cleanup>
111
112 """
Christian Heimes81ee3ef2008-05-04 22:42:01 +0000113 @wraps(func)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000114 def helper(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000115 return _GeneratorContextManager(func, *args, **kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000116 return helper
117
118
Thomas Wouters477c8d52006-05-27 19:21:47 +0000119class closing(object):
120 """Context to automatically close something at the end of a block.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000121
122 Code like this:
123
124 with closing(<module>.open(<arguments>)) as f:
125 <block>
126
127 is equivalent to this:
128
129 f = <module>.open(<arguments>)
130 try:
131 <block>
132 finally:
133 f.close()
134
135 """
Thomas Wouters477c8d52006-05-27 19:21:47 +0000136 def __init__(self, thing):
137 self.thing = thing
138 def __enter__(self):
139 return self.thing
140 def __exit__(self, *exc_info):
141 self.thing.close()