Use the await method properly and just use it instead of some crazy things

jedi/evaluate/context/asynchronous.py

@@ -18,7 +18,8 @@ class AsyncBase(BuiltinOverwrite):
 class Coroutine(AsyncBase):
     special_object_identifier = u'COROUTINE_TYPE'
 
-    def execute_await(self):
+    @publish_method('__await__')
+    def _await(self):
         return self._func_execution_context.get_return_values()
 
 

jedi/evaluate/syntax_tree.py

@@ -91,16 +91,10 @@ def eval_node(context, element):
             context_set = eval_trailer(context, context_set, trailer)
 
         if had_await:
-            await_context_set = ContextSet()
+            await_context_set = context_set.py__getattribute__("__await__")
-            for context in context_set:
+            if not await_context_set:
-                try:
+                debug.warning('Tried to run py__await__ on context %s', context)
-                    func = context.execute_await
+            return await_context_set.execute_evaluated()
-                except AttributeError:
-                    debug.warning('Tried to run execute_await on context %s', context)
-                    pass
-                else:
-                    await_context_set |= func()
-            return await_context_set
         return context_set
     elif typ in ('testlist_star_expr', 'testlist',):
         # The implicit tuple in statements.

test/completion/async_.py

@@ -22,6 +22,8 @@ async def y():
     argh = await x()
     #? int()
     argh
+    #? int()
+    x().__await__()
     return 2
 
 async def asgen():