1 | #!/usr/bin/env python2
|
2 | """Test the with statement for Byterun."""
|
3 |
|
4 | from __future__ import print_function
|
5 |
|
6 | import unittest
|
7 | import six
|
8 |
|
9 | import vmtest
|
10 |
|
11 | PY3 = six.PY3
|
12 |
|
13 | class TestWithStatement(vmtest.VmTestCase):
|
14 |
|
15 | def test_simple_context_manager(self):
|
16 | self.assert_ok("""\
|
17 | class NullContext(object):
|
18 | def __enter__(self):
|
19 | l.append('i')
|
20 | # __enter__ usually returns self, but doesn't have to.
|
21 | return 17
|
22 |
|
23 | def __exit__(self, exc_type, exc_val, exc_tb):
|
24 | l.append('o')
|
25 | return False
|
26 |
|
27 | l = []
|
28 | for i in range(3):
|
29 | with NullContext() as val:
|
30 | assert val == 17
|
31 | l.append('w')
|
32 | l.append('e')
|
33 | l.append('r')
|
34 | s = ''.join(l)
|
35 | print("Look: %r" % s)
|
36 | assert s == "iwoeiwoeiwoer"
|
37 | """)
|
38 |
|
39 | def test_raise_in_context_manager(self):
|
40 | self.assert_ok("""\
|
41 | class NullContext(object):
|
42 | def __enter__(self):
|
43 | l.append('i')
|
44 | return self
|
45 |
|
46 | def __exit__(self, exc_type, exc_val, exc_tb):
|
47 | assert exc_type is ValueError, \\
|
48 | "Expected ValueError: %r" % exc_type
|
49 | l.append('o')
|
50 | return False
|
51 |
|
52 | l = []
|
53 | try:
|
54 | with NullContext():
|
55 | l.append('w')
|
56 | raise ValueError("Boo!")
|
57 | l.append('e')
|
58 | except ValueError:
|
59 | l.append('x')
|
60 | l.append('r')
|
61 | s = ''.join(l)
|
62 | print("Look: %r" % s)
|
63 | assert s == "iwoxr"
|
64 | """)
|
65 |
|
66 | def test_suppressed_raise_in_context_manager(self):
|
67 | self.assert_ok("""\
|
68 | class SuppressingContext(object):
|
69 | def __enter__(self):
|
70 | l.append('i')
|
71 | return self
|
72 |
|
73 | def __exit__(self, exc_type, exc_val, exc_tb):
|
74 | assert exc_type is ValueError, \\
|
75 | "Expected ValueError: %r" % exc_type
|
76 | l.append('o')
|
77 | return True
|
78 |
|
79 | l = []
|
80 | try:
|
81 | with SuppressingContext():
|
82 | l.append('w')
|
83 | raise ValueError("Boo!")
|
84 | l.append('e')
|
85 | except ValueError:
|
86 | l.append('x')
|
87 | l.append('r')
|
88 | s = ''.join(l)
|
89 | print("Look: %r" % s)
|
90 | assert s == "iwoer"
|
91 | """)
|
92 |
|
93 | def test_return_in_with(self):
|
94 | self.assert_ok("""\
|
95 | class NullContext(object):
|
96 | def __enter__(self):
|
97 | l.append('i')
|
98 | return self
|
99 |
|
100 | def __exit__(self, exc_type, exc_val, exc_tb):
|
101 | l.append('o')
|
102 | return False
|
103 |
|
104 | l = []
|
105 | def use_with(val):
|
106 | with NullContext():
|
107 | l.append('w')
|
108 | return val
|
109 | l.append('e')
|
110 |
|
111 | assert use_with(23) == 23
|
112 | l.append('r')
|
113 | s = ''.join(l)
|
114 | print("Look: %r" % s)
|
115 | assert s == "iwor"
|
116 | """)
|
117 |
|
118 | def test_continue_in_with(self):
|
119 | self.assert_ok("""\
|
120 | class NullContext(object):
|
121 | def __enter__(self):
|
122 | l.append('i')
|
123 | return self
|
124 |
|
125 | def __exit__(self, exc_type, exc_val, exc_tb):
|
126 | l.append('o')
|
127 | return False
|
128 |
|
129 | l = []
|
130 | for i in range(3):
|
131 | with NullContext():
|
132 | l.append('w')
|
133 | if i % 2:
|
134 | continue
|
135 | l.append('z')
|
136 | l.append('e')
|
137 |
|
138 | l.append('r')
|
139 | s = ''.join(l)
|
140 | print("Look: %r" % s)
|
141 | assert s == "iwzoeiwoiwzoer"
|
142 | """)
|
143 |
|
144 | def test_break_in_with(self):
|
145 | self.assert_ok("""\
|
146 | class NullContext(object):
|
147 | def __enter__(self):
|
148 | l.append('i')
|
149 | return self
|
150 |
|
151 | def __exit__(self, exc_type, exc_val, exc_tb):
|
152 | l.append('o')
|
153 | return False
|
154 |
|
155 | l = []
|
156 | for i in range(3):
|
157 | with NullContext():
|
158 | l.append('w')
|
159 | if i % 2:
|
160 | break
|
161 | l.append('z')
|
162 | l.append('e')
|
163 |
|
164 | l.append('r')
|
165 | s = ''.join(l)
|
166 | print("Look: %r" % s)
|
167 | assert s == "iwzoeiwor"
|
168 | """)
|
169 |
|
170 | def test_raise_in_with(self):
|
171 | self.assert_ok("""\
|
172 | class NullContext(object):
|
173 | def __enter__(self):
|
174 | l.append('i')
|
175 | return self
|
176 |
|
177 | def __exit__(self, exc_type, exc_val, exc_tb):
|
178 | l.append('o')
|
179 | return False
|
180 |
|
181 | l = []
|
182 | try:
|
183 | with NullContext():
|
184 | l.append('w')
|
185 | raise ValueError("oops")
|
186 | l.append('z')
|
187 | l.append('e')
|
188 | except ValueError as e:
|
189 | assert str(e) == "oops"
|
190 | l.append('x')
|
191 | l.append('r')
|
192 | s = ''.join(l)
|
193 | print("Look: %r" % s)
|
194 | assert s == "iwoxr", "What!?"
|
195 | """)
|
196 |
|
197 | def test_at_context_manager_simplified(self):
|
198 | self.assert_ok("""\
|
199 | class GeneratorContextManager(object):
|
200 | def __init__(self, gen):
|
201 | self.gen = gen
|
202 |
|
203 | def __enter__(self):
|
204 | try:
|
205 | return next(self.gen)
|
206 | except StopIteration:
|
207 | raise RuntimeError("generator didn't yield")
|
208 |
|
209 | def __exit__(self, type, value, traceback):
|
210 | if type is None:
|
211 | try:
|
212 | next(self.gen)
|
213 | except StopIteration:
|
214 | return
|
215 | else:
|
216 | raise RuntimeError("generator didn't stop")
|
217 | else:
|
218 | if value is None:
|
219 | value = type()
|
220 | try:
|
221 | self.gen.throw(type, value, traceback)
|
222 | raise RuntimeError(
|
223 | "generator didn't stop after throw()"
|
224 | )
|
225 | except StopIteration as exc:
|
226 | return exc is not value
|
227 | except:
|
228 | if sys.exc_info()[1] is not value:
|
229 | raise
|
230 |
|
231 | def contextmanager(func):
|
232 | def helper(*args, **kwds):
|
233 | return GeneratorContextManager(func(*args, **kwds))
|
234 | return helper
|
235 |
|
236 | @contextmanager
|
237 | def my_context_manager(val):
|
238 | yield val
|
239 |
|
240 | with my_context_manager(17) as x:
|
241 | assert x == 17
|
242 | """)
|
243 |
|
244 | def test_at_context_manager_complete(self):
|
245 | # The complete code for an @contextmanager example, lifted from
|
246 | # the stdlib.
|
247 | self.assert_ok("""\
|
248 | from _functools import partial
|
249 |
|
250 | WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__doc__')
|
251 | WRAPPER_UPDATES = ('__dict__',)
|
252 |
|
253 | def update_wrapper(wrapper,
|
254 | wrapped,
|
255 | assigned = WRAPPER_ASSIGNMENTS,
|
256 | updated = WRAPPER_UPDATES):
|
257 | for attr in assigned:
|
258 | setattr(wrapper, attr, getattr(wrapped, attr))
|
259 | for attr in updated:
|
260 | getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
|
261 | # Return the wrapper so this can be used as a decorator
|
262 | # via partial().
|
263 | return wrapper
|
264 |
|
265 | def wraps(wrapped,
|
266 | assigned = WRAPPER_ASSIGNMENTS,
|
267 | updated = WRAPPER_UPDATES):
|
268 | return partial(update_wrapper, wrapped=wrapped,
|
269 | assigned=assigned, updated=updated)
|
270 |
|
271 | class GeneratorContextManager(object):
|
272 | def __init__(self, gen):
|
273 | self.gen = gen
|
274 |
|
275 | def __enter__(self):
|
276 | try:
|
277 | return next(self.gen)
|
278 | except StopIteration:
|
279 | raise RuntimeError("generator didn't yield")
|
280 |
|
281 | def __exit__(self, type, value, traceback):
|
282 | if type is None:
|
283 | try:
|
284 | next(self.gen)
|
285 | except StopIteration:
|
286 | return
|
287 | else:
|
288 | raise RuntimeError("generator didn't stop")
|
289 | else:
|
290 | if value is None:
|
291 | value = type()
|
292 | try:
|
293 | self.gen.throw(type, value, traceback)
|
294 | raise RuntimeError(
|
295 | "generator didn't stop after throw()"
|
296 | )
|
297 | except StopIteration as exc:
|
298 | return exc is not value
|
299 | except:
|
300 | if sys.exc_info()[1] is not value:
|
301 | raise
|
302 |
|
303 | def contextmanager(func):
|
304 | @wraps(func)
|
305 | def helper(*args, **kwds):
|
306 | return GeneratorContextManager(func(*args, **kwds))
|
307 | return helper
|
308 |
|
309 | @contextmanager
|
310 | def my_context_manager(val):
|
311 | yield val
|
312 |
|
313 | with my_context_manager(17) as x:
|
314 | assert x == 17
|
315 | """)
|
316 |
|
317 | if PY3:
|
318 | def test_generator_with_context_manager(self):
|
319 | self.assert_ok("""\
|
320 | from contextlib import contextmanager
|
321 |
|
322 | def inner():
|
323 | yield "I'm inner!"
|
324 |
|
325 | def foo():
|
326 | yield from inner()
|
327 |
|
328 | @contextmanager
|
329 | def cmgr():
|
330 | yield "Context Manager!"
|
331 | raise StopIteration(cmgr())
|
332 |
|
333 | def main():
|
334 | with (yield from foo()) as x:
|
335 | print(x)
|
336 |
|
337 | def run(fn, *args):
|
338 | x = fn(*args)
|
339 | while True:
|
340 | try:
|
341 | print(next(x))
|
342 | except StopIteration as e:
|
343 | return e.value
|
344 | run(main)
|
345 | """)
|