[pypy-commit] pypy py3.5-async: Check for illegal return value in GetAwaitableIter after __await__ call, fix GetAwaitableIter call in pyopcodes

[raffael_t]

Author: Raffael Tfirst <raffael.tfirst at gmail.com>
Branch: py3.5-async
Changeset: r85928:60d844fa622a
Date: 2016-07-30 15:32 +0200
http://bitbucket.org/pypy/pypy/changeset/60d844fa622a/

Log:	Check for illegal return value in GetAwaitableIter after __await__
	call, fix GetAwaitableIter call in pyopcodes

diff --git a/pypy/interpreter/generator.py b/pypy/interpreter/generator.py
--- a/pypy/interpreter/generator.py
+++ b/pypy/interpreter/generator.py
@@ -305,13 +305,25 @@
                 block = block.previous
     
     def _GetAwaitableIter(self, space):
-        #check if coroutine
-        if w_iterable.pycode.co_flags & consts.CO_ITERABLE_COROUTINE:
+        #check if generator is a coroutine
+        if self.pycode.co_flags & consts.CO_ITERABLE_COROUTINE:
             return self
-        #look at typeobject.c, change to self.space.lookup(w_manager, "__await__")
-        space.lookup(self, "__await__")
-        res = space.get_and_call_function(w_enter, None)
-        return self
+        w_await = space.lookup(self, "__await__")
+        if w_await is None:
+            raise oefmt(space.w_TypeError,
+                        "object %s can't be used in 'await' expression",
+                        self)
+        res = space.get_and_call_function(w_await, None)
+        if res is not None:
+            if (isinstance(res, Coroutine) or
+                res.pycode.co_flags & consts.CO_ITERABLE_COROUTINE):
+                raise oefmt(space.w_TypeError,
+                            "__await__() returned a coroutine")
+            elif space.lookup(self, "__next__") is None:
+                raise oefmt(space.w_TypeError,
+                        "__await__() returned non-iterator "
+                        "of type '%T'", res)
+        return res
 
 
 class Coroutine(W_Root):
diff --git a/pypy/interpreter/pyopcode.py b/pypy/interpreter/pyopcode.py
--- a/pypy/interpreter/pyopcode.py
+++ b/pypy/interpreter/pyopcode.py
@@ -1442,7 +1442,7 @@
     
     def GET_AWAITABLE(self, oparg, next_instr):
         w_iterable = self.popvalue()
-        w_iter = w_iterable._GetAwaitableIter()
+        w_iter = w_iterable._GetAwaitableIter(self.space)
         self.pushvalue(w_iter)
     
     def SETUP_ASYNC_WITH(self, offsettoend, next_instr):
@@ -1468,13 +1468,13 @@
     def GET_AITER(self, oparg, next_instr):
         w_obj = self.popvalue()
         w_iter = w_obj.aiter()
-        w_awaitable = w_iter._GetAwaitableIter()
+        w_awaitable = w_iter._GetAwaitableIter(self.space)
         self.pushvalue(w_awaitable)
     
     def GET_ANEXT(self, oparg, next_instr):
         w_aiter = self.popvalue()
         w_next_iter = w_aiter.anext()
-        w_awaitable = w_next_iter._GetAwaitableIter()
+        w_awaitable = w_next_iter._GetAwaitableIter(self.space)
         self.pushvalue(w_awaitable)
         
 ### ____________________________________________________________ ###