Replace AsyncGenerator

jedi/evaluate/__init__.py

@@ -134,7 +134,9 @@ class Evaluator(object):
             self,
         )
 
-    def import_module(self, import_names, parent_module_context, sys_path):
+    def import_module(self, import_names, parent_module_context=None, sys_path=None):
+        if sys_path is None:
+            sys_path = self.get_sys_path()
         try:
             return self.module_cache.get(import_names)
         except KeyError:
@@ -150,13 +152,15 @@ class Evaluator(object):
         module_name = 'builtins'
         if self.environment.version_info.major == 2:
             module_name = '__builtin__'
-        builtins_module, = self.import_module(
-            (module_name,),
-            parent_module_context=None,
-            sys_path=()
-        )
+        builtins_module, = self.import_module((module_name,), sys_path=())
         return builtins_module
 
+    @property
+    @evaluator_function_cache()
+    def typing_module(self):
+        typing_module, = self.import_module(('typing',))
+        return typing_module
+
     def reset_recursion_limitations(self):
         self.recursion_detector = recursion.RecursionDetector()
         self.execution_recursion_detector = recursion.ExecutionRecursionDetector(self)

jedi/evaluate/base_context.py

@@ -31,9 +31,10 @@ class HelperContextMixin:
     def execute_evaluated(self, *value_list):
         return execute_evaluated(self, *value_list)
 
-    def merge_types_of_iterate(self):
+    def merge_types_of_iterate(self, contextualized_node=None, is_async=False):
         return ContextSet.from_sets(
-            lazy_context.infer() for lazy_context in self.iterate()
+            lazy_context.infer()
+            for lazy_context in self.iterate(contextualized_node, is_async)
         )
 
     @Python3Method
@@ -81,6 +82,16 @@ class Context(HelperContextMixin, BaseContext):
 
     def iterate(self, contextualized_node=None, is_async=False):
         debug.dbg('iterate %s', self)
+        if is_async:
+            from jedi.evaluate.lazy_context import LazyKnownContexts
+            return iter([
+                LazyKnownContexts(
+                    self.py__getattribute__('__aiter__').execute_evaluated()
+                        .py__getattribute__('__anext__').execute_evaluated()
+                        .py__getattribute__('__await__').execute_evaluated()
+                        .py__stop_iteration_returns()
+                )  # noqa
+            ])
         try:
             if is_async:
                 iter_method = self.py__aiter__
@@ -140,6 +151,10 @@ class Context(HelperContextMixin, BaseContext):
                 return clean_scope_docstring(self.tree_node)
         return None
 
+    def py__stop_iteration_returns(self):
+        debug.warning("Not possible to return the stop iterations of %s", self)
+        return NO_CONTEXTS
+
 
 def iterate_contexts(contexts, contextualized_node=None, is_async=False):
     """

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()