Use async generator/async functions from typeshed

jedi/evaluate/__init__.py

@@ -159,7 +159,7 @@ class Evaluator(object):
     @evaluator_function_cache()
     def typing_module(self):
         typing_module, = self.import_module(('typing',))
-        return typing_module
+        return typing_module.stub_context
 
     def reset_recursion_limitations(self):
         self.recursion_detector = recursion.RecursionDetector()

jedi/evaluate/compiled/access.py

@@ -474,9 +474,6 @@ class _SPECIAL_OBJECTS(object):
     MODULE_CLASS = types.ModuleType
     GENERATOR_OBJECT = _a_generator(1.0)
     BUILTINS = builtins
-    COROUTINE = _coroutine
-    COROUTINE_WRAPPER = _coroutine_wrapper
-    ASYNC_GENERATOR = _async_generator
 
 
 def get_special_object(evaluator, identifier):

jedi/evaluate/context/asynchronous.py

@@ -1,30 +0,0 @@
-from jedi.evaluate.filters import publish_method, BuiltinOverwrite
-from jedi.evaluate.base_context import ContextSet
-
-
-class AsyncBase(BuiltinOverwrite):
-    def __init__(self, evaluator, func_execution_context):
-        super(AsyncBase, self).__init__(evaluator)
-        self.func_execution_context = func_execution_context
-
-    @property
-    def name(self):
-        return self.get_object().name
-
-    def __repr__(self):
-        return "<%s of %s>" % (type(self).__name__, self.func_execution_context)
-
-
-class Coroutine(AsyncBase):
-    special_object_identifier = u'COROUTINE'
-
-    @publish_method('__await__')
-    def _await(self):
-        return ContextSet(CoroutineWrapper(self.evaluator, self.func_execution_context))
-
-
-class CoroutineWrapper(AsyncBase):
-    special_object_identifier = u'COROUTINE_WRAPPER'
-
-    def py__stop_iteration_returns(self):
-        return self.func_execution_context.get_return_values()

jedi/evaluate/context/function.py

@@ -17,7 +17,6 @@ from jedi.evaluate.base_context import ContextualizedNode, NO_CONTEXTS, \
 from jedi.evaluate.lazy_context import LazyKnownContexts, LazyKnownContext, \
     LazyTreeContext
 from jedi.evaluate.context import iterable
-from jedi.evaluate.context import asynchronous
 from jedi import parser_utils
 from jedi.evaluate.parser_cache import get_yield_exprs
 
@@ -312,22 +311,35 @@ class FunctionExecutionContext(TreeContext):
             if is_generator:
                 if evaluator.environment.version_info < (3, 6):
                     return NO_CONTEXTS
-                async_generator_classes = evaluator.typing_module \
+                async_generator_class, = evaluator.typing_module \
                     .py__getattribute__('AsyncGenerator')
+
+                c = async_generator_class.stub_context
                 yield_contexts = self.merge_yield_contexts(is_async=True)
-                return ContextSet.from_iterable(
+                return ContextSet(
                     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
+                    )
-                )
+                ).execute_annotation()
             else:
                 if evaluator.environment.version_info < (3, 5):
                     return NO_CONTEXTS
-                return ContextSet(asynchronous.Coroutine(evaluator, self))
+                async_class, = evaluator.typing_module .py__getattribute__('Coroutine')
+                return_contexts = self.get_return_values()
+                c = async_class.stub_context
+                return ContextSet(
+                    AnnotatedSubClass(
+                        evaluator,
+                        parent_context=c.parent_context,
+                        tree_node=c.tree_node,
+                        # Only the first generic is relevant.
+                        given_types=(return_contexts.py__class__(), NO_CONTEXTS, NO_CONTEXTS)
+                    )
+                ).execute_annotation()
         else:
             if is_generator:
                 return ContextSet(iterable.Generator(evaluator, self))