py/compile: Fix async for's stack handling of iterator expression.

Prior to this fix, async for assumed the iterator expression was a simple
identifier, and used that identifier as a local to store the intermediate
iterator object.  This is incorrect behaviour.

This commit fixes the issue by keeping the iterator object on the stack as
an anonymous local variable.

Fixes issue #11511.

Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 58 insertions, 12 deletions

diff --git a/tests/basics/async_for.py b/tests/basics/async_for.py
index 5fd0540828..f54f70238c 100644
--- a/tests/basics/async_for.py
+++ b/tests/basics/async_for.py
@@ -1,29 +1,75 @@
 # test basic async for execution
 # example taken from PEP0492
 
+
 class AsyncIteratorWrapper:
     def __init__(self, obj):
-        print('init')
-        self._it = iter(obj)
+        print("init")
+        self._obj = obj
+
+    def __repr__(self):
+        return "AsyncIteratorWrapper-" + self._obj
 
     def __aiter__(self):
-        print('aiter')
-        return self
+        print("aiter")
+        return AsyncIteratorWrapperIterator(self._obj)
+
+
+class AsyncIteratorWrapperIterator:
+    def __init__(self, obj):
+        print("init")
+        self._it = iter(obj)
 
     async def __anext__(self):
-        print('anext')
+        print("anext")
         try:
             value = next(self._it)
         except StopIteration:
             raise StopAsyncIteration
         return value
 
-async def coro():
-    async for letter in AsyncIteratorWrapper('abc'):
+
+def run_coro(c):
+    print("== start ==")
+    try:
+        c.send(None)
+    except StopIteration:
+        print("== finish ==")
+
+
+async def coro0():
+    async for letter in AsyncIteratorWrapper("abc"):
         print(letter)
 
-o = coro()
-try:
-    o.send(None)
-except StopIteration:
-    print('finished')
+
+run_coro(coro0())
+
+
+async def coro1():
+    a = AsyncIteratorWrapper("def")
+    async for letter in a:
+        print(letter)
+    print(a)
+
+
+run_coro(coro1())
+
+a_global = AsyncIteratorWrapper("ghi")
+
+
+async def coro2():
+    async for letter in a_global:
+        print(letter)
+    print(a_global)
+
+
+run_coro(coro2())
+
+
+async def coro3(a):
+    async for letter in a:
+        print(letter)
+    print(a)
+
+
+run_coro(coro3(AsyncIteratorWrapper("jkl")))