Preparations for some async changes

jedi/evaluate/context/function.py

@@ -13,7 +13,7 @@ from jedi.evaluate.arguments import AnonymousArguments
 from jedi.evaluate.filters import ParserTreeFilter, FunctionExecutionFilter
 from jedi.evaluate.names import ContextName, AbstractNameDefinition, ParamName
 from jedi.evaluate.base_context import ContextualizedNode, NO_CONTEXTS, \
-    ContextSet, TreeContext, ContextWrapper
+    ContextSet, TreeContext, ContextWrapper, LazyContextWrapper
 from jedi.evaluate.lazy_context import LazyKnownContexts, LazyKnownContext, \
     LazyTreeContext
 from jedi.evaluate.context import iterable
@@ -353,6 +353,7 @@ class FunctionExecutionContext(TreeContext):
             else:
                 if evaluator.environment.version_info < (3, 5):
                     return NO_CONTEXTS
+                #return ContextSet({CoroutineObject(evaluator, self)})
                 async_classes = evaluator.typing_module.py__getattribute__('Coroutine')
                 return_contexts = self.get_return_values()
                 # Only the first generic is relevant.
@@ -367,6 +368,34 @@ class FunctionExecutionContext(TreeContext):
                 return self.get_return_values()
 
 
+class CoroutineObject(LazyContextWrapper):
+    def __init__(self, evaluator, function_execution):
+        self.evaluator = evaluator
+        self._function_execution = function_execution
+
+    def _get_wrapped_context(self):
+        c, = self.evaluator.typing_module.py__getattribute__('Coroutine') \
+            .execute_annotation()
+        return c
+
+    def py__await__(self):
+        return ContextSet({CoroutineWrapper(self.evaluator, self._function_execution)})
+
+
+class CoroutineWrapper(LazyContextWrapper):
+    def __init__(self, evaluator, function_execution):
+        self.evaluator = evaluator
+        self._function_execution = function_execution
+
+    def _get_wrapped_context(self):
+        c, = self.evaluator.typing_module.py__getattribute__('Generator') \
+            .execute_annotation()
+        return c
+
+    def py__stop_iteration_returns(self):
+        return self._function_execution.get_return_values()
+
+
 class OverloadedFunctionContext(FunctionMixin, ContextWrapper):
     def __init__(self, function, overloaded_functions):
         super(OverloadedFunctionContext, self).__init__(function)