Issue 24316: Wrap gen objects returned from callables in types.coroutine
This commit is contained in:
parent
c565cd5d1b
commit
13f7723d81
|
@ -1206,28 +1206,51 @@ class CoroutineTests(unittest.TestCase):
|
|||
@types.coroutine
|
||||
def foo():
|
||||
pass
|
||||
@types.coroutine
|
||||
def gen():
|
||||
def _gen(): yield
|
||||
return _gen()
|
||||
|
||||
for sample in (foo, gen):
|
||||
with self.assertRaisesRegex(TypeError,
|
||||
'callable wrapped .* non-coroutine'):
|
||||
sample()
|
||||
with self.assertRaisesRegex(TypeError,
|
||||
'callable wrapped .* non-coroutine'):
|
||||
foo()
|
||||
|
||||
def test_duck_coro(self):
|
||||
class CoroLike:
|
||||
def send(self): pass
|
||||
def throw(self): pass
|
||||
def close(self): pass
|
||||
def __await__(self): pass
|
||||
def __await__(self): return self
|
||||
|
||||
coro = CoroLike()
|
||||
@types.coroutine
|
||||
def foo():
|
||||
return coro
|
||||
self.assertIs(coro, foo())
|
||||
self.assertIs(foo().__await__(), coro)
|
||||
|
||||
def test_duck_gen(self):
|
||||
class GenLike:
|
||||
def send(self): pass
|
||||
def throw(self): pass
|
||||
def close(self): pass
|
||||
def __iter__(self): return self
|
||||
def __next__(self): pass
|
||||
|
||||
gen = GenLike()
|
||||
@types.coroutine
|
||||
def foo():
|
||||
return gen
|
||||
self.assertIs(foo().__await__(), gen)
|
||||
|
||||
with self.assertRaises(AttributeError):
|
||||
foo().gi_code
|
||||
|
||||
def test_gen(self):
|
||||
def gen(): yield
|
||||
gen = gen()
|
||||
@types.coroutine
|
||||
def foo(): return gen
|
||||
self.assertIs(foo().__await__(), gen)
|
||||
|
||||
for name in ('__name__', '__qualname__', 'gi_code',
|
||||
'gi_running', 'gi_frame'):
|
||||
self.assertIs(getattr(foo(), name),
|
||||
getattr(gen, name))
|
||||
|
||||
def test_genfunc(self):
|
||||
def gen():
|
||||
|
|
40
Lib/types.py
40
Lib/types.py
|
@ -166,32 +166,64 @@ def coroutine(func):
|
|||
|
||||
# We don't want to import 'dis' or 'inspect' just for
|
||||
# these constants.
|
||||
_CO_GENERATOR = 0x20
|
||||
_CO_ITERABLE_COROUTINE = 0x100
|
||||
CO_GENERATOR = 0x20
|
||||
CO_ITERABLE_COROUTINE = 0x100
|
||||
|
||||
if not callable(func):
|
||||
raise TypeError('types.coroutine() expects a callable')
|
||||
|
||||
if (isinstance(func, FunctionType) and
|
||||
isinstance(getattr(func, '__code__', None), CodeType) and
|
||||
(func.__code__.co_flags & _CO_GENERATOR)):
|
||||
(func.__code__.co_flags & CO_GENERATOR)):
|
||||
|
||||
# TODO: Implement this in C.
|
||||
co = func.__code__
|
||||
func.__code__ = CodeType(
|
||||
co.co_argcount, co.co_kwonlyargcount, co.co_nlocals,
|
||||
co.co_stacksize,
|
||||
co.co_flags | _CO_ITERABLE_COROUTINE,
|
||||
co.co_flags | CO_ITERABLE_COROUTINE,
|
||||
co.co_code,
|
||||
co.co_consts, co.co_names, co.co_varnames, co.co_filename,
|
||||
co.co_name, co.co_firstlineno, co.co_lnotab, co.co_freevars,
|
||||
co.co_cellvars)
|
||||
return func
|
||||
|
||||
# The following code is primarily to support functions that
|
||||
# return generator-like objects (for instance generators
|
||||
# compiled with Cython).
|
||||
|
||||
class GeneratorWrapper:
|
||||
def __init__(self, gen):
|
||||
self.__wrapped__ = gen
|
||||
self.send = gen.send
|
||||
self.throw = gen.throw
|
||||
self.close = gen.close
|
||||
self.__name__ = getattr(gen, '__name__', None)
|
||||
self.__qualname__ = getattr(gen, '__qualname__', None)
|
||||
@property
|
||||
def gi_code(self):
|
||||
return self.__wrapped__.gi_code
|
||||
@property
|
||||
def gi_frame(self):
|
||||
return self.__wrapped__.gi_frame
|
||||
@property
|
||||
def gi_running(self):
|
||||
return self.__wrapped__.gi_running
|
||||
def __next__(self):
|
||||
return next(self.__wrapped__)
|
||||
def __iter__(self):
|
||||
return self.__wrapped__
|
||||
__await__ = __iter__
|
||||
|
||||
@_functools.wraps(func)
|
||||
def wrapped(*args, **kwargs):
|
||||
coro = func(*args, **kwargs)
|
||||
if coro.__class__ is GeneratorType:
|
||||
return GeneratorWrapper(coro)
|
||||
# slow checks
|
||||
if not isinstance(coro, _collections_abc.Coroutine):
|
||||
if isinstance(coro, _collections_abc.Generator):
|
||||
return GeneratorWrapper(coro)
|
||||
raise TypeError(
|
||||
'callable wrapped with types.coroutine() returned '
|
||||
'non-coroutine: {!r}'.format(coro))
|
||||
|
|
Loading…
Reference in New Issue