Source

pypy / pypy / interpreter / function.py

Diff from to

pypy/interpreter/function.py

             w_arg = frame.peekvalue(nargs-1-i)
             new_frame.locals_stack_w[i] = w_arg
 
-        return new_frame.run()
+        result = new_frame.run()
+        result = self._wrap_await_result(result)
+        return result
+
+    def _wrap_await_result(self, result):
+        """
+        Python-extension for the 'await' keyword. Wrap the last iteration of an
+        await-generator through executioncontext.w_await_schedular_func.
+        """
+        # Does the code object of this function contain the CO_AWAIT_GENERATOR
+        # flag? (Set by the ast.)
+        from pypy.interpreter.pycode import PyCode
+        from pypy.interpreter.astcompiler.consts import CO_GENERATOR, CO_AWAIT_GENERATOR
+        is_await_generator = (self.getcode().co_flags & CO_AWAIT_GENERATOR)
+
+        # If so, and if an I/O schedular is set, get the result through the
+        # I/O schedular function.
+        if is_await_generator:
+            context = self.space.getexecutioncontext()
+            if context.w_await_schedular_func:
+                schedular = context.w_await_schedular_func
+                assert isinstance(schedular, Function)
+                result = schedular.funccall(result)
+            else:
+                raise OperationError(self.space.w_TypeError,
+                    self.space.wrap("Calling an asyncronous function (with an 'await'-keyword) while no I/O schedular has been set")
+                )
+        return result
 
     @jit.unroll_safe
     def _flat_pycall_defaults(self, code, nargs, frame, defs_to_load):