Added type checking of async generators via @TypeChecked

CHANGELOG.rst

@@ -5,14 +5,14 @@ This library adheres to `Semantic Versioning 2.0 <https://semver.org/#semantic-v
 
 **UNRELEASED**
 
-- Added generator yield type checking in ``TypeChecker`` for regular generators
+- Added yield type checking via ``TypeChecker`` for regular generators
+- Added yield, send and return type checking via ``@typechecked`` for regular and async generators
 - Silenced ``TypeChecker`` warnings about async generators
 - Fixed bogus ``TypeError`` on ``Type[Any]``
 - Fixed bogus ``TypeChecker`` warnings when an exception is raised from a type checked function
 - Accept a ``bytearray`` where ``bytes`` are expected, as per `python/typing#552`_
 - Added policies for dealing with unmatched forward references
 - Added support for using ``@typechecked`` as a class decorator
-- Added type checking for generators via ``@typechecked``
 
 .. _python/typing#552: https://github.com/python/typing/issues/552
 

README.rst

@@ -22,6 +22,7 @@ There are three principal ways to use type checking, each with its pros and cons
    * 100% reliable at finding the function object to be checked (does not need to check the garbage
      collector)
    * can check the type of the return value
+   * wraps returned generators (async or regular) and type checks yields, sends and returns
    * adds an extra frame to the call stack for every call to a decorated function
 #. using ``with TypeChecker('packagename'):``:
 

tests/test_typeguard_py36.py

@@ -2,7 +2,59 @@
 
 import pytest
 
-from typeguard import TypeChecker
+from typeguard import TypeChecker, typechecked
+
+
+class TestTypeChecked:
+    def test_async_generator(self):
+        async def run_generator():
+            @typechecked
+            async def genfunc() -> AsyncGenerator[int, str]:
+                values.append((yield 2))
+                values.append((yield 3))
+                values.append((yield 4))
+
+            gen = genfunc()
+
+            value = await gen.asend(None)
+            with pytest.raises(StopAsyncIteration):
+                while True:
+                    value = await gen.asend(str(value))
+                    assert isinstance(value, int)
+
+        values = []
+        coro = run_generator()
+        try:
+            for elem in coro.__await__():
+                print(elem)
+        except StopAsyncIteration as exc:
+            values = exc.value
+
+        assert values == ['2', '3', '4']
+
+    def test_async_generator_bad_yield(self):
+        @typechecked
+        async def genfunc() -> AsyncGenerator[int, str]:
+            yield 'foo'
+
+        gen = genfunc()
+        with pytest.raises(TypeError) as exc:
+            next(gen.__anext__().__await__())
+
+        exc.match('type of value yielded from generator must be int; got str instead')
+
+    def test_async_generator_bad_send(self):
+        @typechecked
+        async def genfunc() -> AsyncGenerator[int, str]:
+            yield 1
+            yield 2
+
+        gen = genfunc()
+        pytest.raises(StopIteration, next, gen.__anext__().__await__())
+        with pytest.raises(TypeError) as exc:
+            next(gen.asend(2).__await__())
+
+        exc.match('type of value sent to generator must be str; got int instead')
 
 
 class TestTypeChecker:

typeguard/__init__.py

@@ -604,6 +604,33 @@ def send(self, obj):
         return value
 
 
+class TypeCheckedAsyncGenerator:
+    def __init__(self, wrapped: AsyncGenerator, memo: _CallMemo):
+        self.__wrapped = wrapped
+        self.__memo = memo
+        self.__yield_type, self.__send_type = memo.type_hints['return'].__args__
+        self.__initialized = False
+
+    async def __aiter__(self):
+        return self
+
+    def __anext__(self):
+        return self.asend(None)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.__wrapped, name)
+
+    async def asend(self, obj):
+        if self.__initialized:
+            check_type('value sent to generator', obj, self.__send_type, memo=self.__memo)
+        else:
+            self.__initialized = True
+
+        value = await self.__wrapped.asend(obj)
+        check_type('value yielded from generator', value, self.__yield_type, memo=self.__memo)
+        return value
+
+
 @overload
 def typechecked(*, always: bool = False) -> Callable[[T_CallableOrType], T_CallableOrType]:
     ...
@@ -667,9 +694,18 @@ async def async_wrapper(*args, **kwargs):
         check_return_type(retval, memo)
         return retval
 
+    def asyncgen_wrapper(*args, **kwargs):
+        memo = _CallMemo(func, args=args, kwargs=kwargs)
+        check_argument_types(memo)
+        retval = func(*args, **kwargs)
+        return TypeCheckedAsyncGenerator(retval, memo)
+
     if inspect.iscoroutinefunction(func):
         if func.__code__ is not async_wrapper.__code__:
             return wraps(func)(async_wrapper)
+    elif isasyncgenfunction(func):
+        if func.__code__ is not asyncgen_wrapper.__code__:
+            return wraps(func)(asyncgen_wrapper)
     else:
         if func.__code__ is not wrapper.__code__:
             return wraps(func)(wrapper)