Replace AsyncGenerator

jedi/evaluate/context/asynchronous.py

@@ -28,11 +28,3 @@ class CoroutineWrapper(AsyncBase):
 
     def py__stop_iteration_returns(self):
         return self.func_execution_context.get_return_values()
-
-
-class AsyncGenerator(AsyncBase):
-    """Handling of `yield` functions."""
-    special_object_identifier = u'ASYNC_GENERATOR'
-
-    def py__aiter__(self):
-        return self.func_execution_context.get_yield_lazy_contexts(is_async=True)

jedi/evaluate/context/function.py

@@ -267,6 +267,12 @@ class FunctionExecutionContext(TreeContext):
                             for result in self._get_yield_lazy_context(yield_in_same_for_stmt):
                                 yield result
 
+    def merge_yield_contexts(self, is_async=False):
+        return ContextSet.from_sets(
+            lazy_context.infer()
+            for lazy_context in self.get_yield_lazy_contexts()
+        )
+
     def get_filters(self, search_global=False, until_position=None, origin_scope=None):
         yield self.function_execution_filter(self.evaluator, self,
                                              until_position=until_position,
@@ -300,12 +306,24 @@ class FunctionExecutionContext(TreeContext):
         evaluator = self.evaluator
         is_coroutine = self.tree_node.parent.type == 'async_stmt'
         is_generator = bool(get_yield_exprs(evaluator, self.tree_node))
+        from jedi.evaluate.context.typing import AnnotatedSubClass
 
         if is_coroutine:
             if is_generator:
                 if evaluator.environment.version_info < (3, 6):
                     return NO_CONTEXTS
-                return ContextSet(asynchronous.AsyncGenerator(evaluator, self))
+                async_generator_classes = evaluator.typing_module \
+                    .py__getattribute__('AsyncGenerator')
+                yield_contexts = self.merge_yield_contexts(is_async=True)
+                return ContextSet.from_iterable(
+                    AnnotatedSubClass(
+                        evaluator,
+                        parent_context=c.parent_context,
+                        tree_node=c.tree_node,
+                        # The contravariant doesn't seem to be defined.
+                        given_types=(yield_contexts.py__class__(), NO_CONTEXTS)
+                    ) for c in async_generator_classes
+                )
             else:
                 if evaluator.environment.version_info < (3, 5):
                     return NO_CONTEXTS

jedi/evaluate/context/typing.py

@@ -8,7 +8,7 @@ from jedi import debug
 from jedi.evaluate.cache import evaluator_method_cache
 from jedi.evaluate.compiled import builtin_from_name, CompiledObject
 from jedi.evaluate.base_context import ContextSet, NO_CONTEXTS, Context, \
-    iterator_to_context_set, HelperContextMixin
+    iterator_to_context_set, HelperContextMixin, ContextWrapper
 from jedi.evaluate.lazy_context import LazyKnownContexts, LazyKnownContext
 from jedi.evaluate.context.iterable import SequenceLiteralContext
 from jedi.evaluate.arguments import repack_with_argument_clinic
@@ -528,6 +528,10 @@ class _AbstractAnnotatedClass(ClassContext):
             ) for class_set1, class_set2 in zip(given_params1, given_params2)
         )
 
+    def py__call__(self, arguments):
+        instance, = super(_AbstractAnnotatedClass, self).py__call__(arguments)
+        return ContextSet(InstanceWrapper(instance))
+
     def get_given_types(self):
         raise NotImplementedError
 
@@ -599,3 +603,15 @@ class LazyAnnotatedBaseClass(object):
                     # case just add it to the context set.
                     new |= ContextSet(type_var)
             yield new
+
+
+class InstanceWrapper(ContextWrapper):
+    def py__stop_iteration_returns(self):
+        cls = self._wrapped_context.class_context
+        if cls.py__name__() == 'Generator':
+            given_types = cls.get_given_types()
+            try:
+                return given_types[2].execute_annotation()
+            except IndexError:
+                pass
+        return self._wrapped_context.py__stop_iteration_returns()