blob: 2f8f00d2964c241d58fe7de53ee083f88f4d209d [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
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00005
Michael Foordb3a89842010-06-30 12:17:50 +00006__all__ = ["contextmanager", "closing", "ContextDecorator"]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00007
Michael Foordb3a89842010-06-30 12:17:50 +00008
9class ContextDecorator(object):
10 "A base class or mixin that enables context managers to work as decorators."
Nick Coghlan0ded3e32011-05-05 23:49:25 +100011
12 def _recreate_cm(self):
13 """Return a recreated instance of self.
Nick Coghlanfdc2c552011-05-06 00:02:12 +100014
Nick Coghlan0ded3e32011-05-05 23:49:25 +100015 Allows otherwise one-shot context managers like
16 _GeneratorContextManager to support use as
17 decorators via implicit recreation.
Nick Coghlanfdc2c552011-05-06 00:02:12 +100018
Nick Coghlan0ded3e32011-05-05 23:49:25 +100019 Note: this is a private interface just for _GCM in 3.2 but will be
20 renamed and documented for third party use in 3.3
21 """
22 return self
23
Michael Foordb3a89842010-06-30 12:17:50 +000024 def __call__(self, func):
25 @wraps(func)
26 def inner(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +100027 with self._recreate_cm():
Michael Foordb3a89842010-06-30 12:17:50 +000028 return func(*args, **kwds)
29 return inner
30
31
Antoine Pitrou67b212e2011-01-08 09:55:31 +000032class _GeneratorContextManager(ContextDecorator):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000033 """Helper for @contextmanager decorator."""
34
Nick Coghlan0ded3e32011-05-05 23:49:25 +100035 def __init__(self, func, *args, **kwds):
36 self.gen = func(*args, **kwds)
37 self.func, self.args, self.kwds = func, args, kwds
38
39 def _recreate_cm(self):
40 # _GCM instances are one-shot context managers, so the
41 # CM must be recreated each time a decorated function is
42 # called
43 return self.__class__(self.func, *self.args, **self.kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000044
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000045 def __enter__(self):
46 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000047 return next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000048 except StopIteration:
49 raise RuntimeError("generator didn't yield")
50
51 def __exit__(self, type, value, traceback):
52 if type is None:
53 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000054 next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000055 except StopIteration:
56 return
57 else:
58 raise RuntimeError("generator didn't stop")
59 else:
Guido van Rossum2cc30da2007-11-02 23:46:40 +000060 if value is None:
61 # Need to force instantiation so we can reliably
62 # tell if we get the same exception back
63 value = type()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000064 try:
65 self.gen.throw(type, value, traceback)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000066 raise RuntimeError("generator didn't stop after throw()")
Guido van Rossumb940e112007-01-10 16:19:56 +000067 except StopIteration as exc:
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000068 # Suppress the exception *unless* it's the same exception that
69 # was passed to throw(). This prevents a StopIteration
70 # raised inside the "with" statement from being suppressed
71 return exc is not value
72 except:
73 # only re-raise if it's *not* the exception that was
74 # passed to throw(), because __exit__() must not raise
75 # an exception unless __exit__() itself failed. But throw()
76 # has to raise the exception to signal propagation, so this
77 # fixes the impedance mismatch between the throw() protocol
78 # and the __exit__() protocol.
79 #
80 if sys.exc_info()[1] is not value:
81 raise
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000082
83
84def contextmanager(func):
85 """@contextmanager decorator.
86
87 Typical usage:
88
89 @contextmanager
90 def some_generator(<arguments>):
91 <setup>
92 try:
93 yield <value>
94 finally:
95 <cleanup>
96
97 This makes this:
98
99 with some_generator(<arguments>) as <variable>:
100 <body>
101
102 equivalent to this:
103
104 <setup>
105 try:
106 <variable> = <value>
107 <body>
108 finally:
109 <cleanup>
110
111 """
Christian Heimes81ee3ef2008-05-04 22:42:01 +0000112 @wraps(func)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000113 def helper(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000114 return _GeneratorContextManager(func, *args, **kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000115 return helper
116
117
Thomas Wouters477c8d52006-05-27 19:21:47 +0000118class closing(object):
119 """Context to automatically close something at the end of a block.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000120
121 Code like this:
122
123 with closing(<module>.open(<arguments>)) as f:
124 <block>
125
126 is equivalent to this:
127
128 f = <module>.open(<arguments>)
129 try:
130 <block>
131 finally:
132 f.close()
133
134 """
Thomas Wouters477c8d52006-05-27 19:21:47 +0000135 def __init__(self, thing):
136 self.thing = thing
137 def __enter__(self):
138 return self.thing
139 def __exit__(self, *exc_info):
140 self.thing.close()