mirror of
https://github.com/python/cpython.git
synced 2025-10-26 03:04:41 +00:00
33 lines
928 B
Python
33 lines
928 B
Python
class GeneratorContextManager(object):
|
|
def __init__(self, gen):
|
|
self.gen = gen
|
|
|
|
def __context__(self):
|
|
return self
|
|
|
|
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:
|
|
try:
|
|
self.gen.throw(type, value, traceback)
|
|
except (type, StopIteration):
|
|
return
|
|
else:
|
|
raise RuntimeError("generator caught exception")
|
|
|
|
def contextmanager(func):
|
|
def helper(*args, **kwds):
|
|
return GeneratorContextManager(func(*args, **kwds))
|
|
return helper
|