diff --git a/jedi/evaluate/context/asynchronous.py b/jedi/evaluate/context/asynchronous.py index cc8dab27..1061ebbf 100644 --- a/jedi/evaluate/context/asynchronous.py +++ b/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() diff --git a/jedi/evaluate/syntax_tree.py b/jedi/evaluate/syntax_tree.py index ae781c2a..122e3cf7 100644 --- a/jedi/evaluate/syntax_tree.py +++ b/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() - for context in context_set: - try: - func = context.execute_await - except AttributeError: - debug.warning('Tried to run execute_await on context %s', context) - pass - else: - await_context_set |= func() - return await_context_set + await_context_set = context_set.py__getattribute__("__await__") + if not await_context_set: + debug.warning('Tried to run py__await__ on context %s', context) + return await_context_set.execute_evaluated() return context_set elif typ in ('testlist_star_expr', 'testlist',): # The implicit tuple in statements. diff --git a/test/completion/async_.py b/test/completion/async_.py index 02e35cc6..95a98853 100644 --- a/test/completion/async_.py +++ b/test/completion/async_.py @@ -22,6 +22,8 @@ async def y(): argh = await x() #? int() argh + #? int() + x().__await__() return 2 async def asgen():