From 41782e497092a27efbae20208ce7d48c657e74cb Mon Sep 17 00:00:00 2001
From: Yury Selivanov <yury@magic.io>
Date: Wed, 16 Nov 2016 18:16:17 -0500
Subject: Issue #28721: Fix asynchronous generators aclose() and athrow()

---
 Lib/test/test_asyncgen.py | 140 ++++++++++++++++++++++++++++++++++++++++++++++
 Misc/NEWS                 |   3 +
 Objects/genobject.c       |  14 ++++-
 3 files changed, 154 insertions(+), 3 deletions(-)

diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py
index 68d2029..34ab8a0 100644
--- a/Lib/test/test_asyncgen.py
+++ b/Lib/test/test_asyncgen.py
@@ -450,6 +450,41 @@ class AsyncGenAsyncioTest(unittest.TestCase):
 
         self.loop.run_until_complete(run())
 
+    def test_async_gen_asyncio_anext_06(self):
+        DONE = 0
+
+        # test synchronous generators
+        def foo():
+            try:
+                yield
+            except:
+                pass
+        g = foo()
+        g.send(None)
+        with self.assertRaises(StopIteration):
+            g.send(None)
+
+        # now with asynchronous generators
+
+        async def gen():
+            nonlocal DONE
+            try:
+                yield
+            except:
+                pass
+            DONE = 1
+
+        async def run():
+            nonlocal DONE
+            g = gen()
+            await g.asend(None)
+            with self.assertRaises(StopAsyncIteration):
+                await g.asend(None)
+            DONE += 10
+
+        self.loop.run_until_complete(run())
+        self.assertEqual(DONE, 11)
+
     def test_async_gen_asyncio_anext_tuple(self):
         async def foo():
             try:
@@ -594,6 +629,76 @@ class AsyncGenAsyncioTest(unittest.TestCase):
         self.loop.run_until_complete(run())
         self.assertEqual(DONE, 1)
 
+    def test_async_gen_asyncio_aclose_10(self):
+        DONE = 0
+
+        # test synchronous generators
+        def foo():
+            try:
+                yield
+            except:
+                pass
+        g = foo()
+        g.send(None)
+        g.close()
+
+        # now with asynchronous generators
+
+        async def gen():
+            nonlocal DONE
+            try:
+                yield
+            except:
+                pass
+            DONE = 1
+
+        async def run():
+            nonlocal DONE
+            g = gen()
+            await g.asend(None)
+            await g.aclose()
+            DONE += 10
+
+        self.loop.run_until_complete(run())
+        self.assertEqual(DONE, 11)
+
+    def test_async_gen_asyncio_aclose_11(self):
+        DONE = 0
+
+        # test synchronous generators
+        def foo():
+            try:
+                yield
+            except:
+                pass
+            yield
+        g = foo()
+        g.send(None)
+        with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'):
+            g.close()
+
+        # now with asynchronous generators
+
+        async def gen():
+            nonlocal DONE
+            try:
+                yield
+            except:
+                pass
+            yield
+            DONE += 1
+
+        async def run():
+            nonlocal DONE
+            g = gen()
+            await g.asend(None)
+            with self.assertRaisesRegex(RuntimeError, 'ignored GeneratorExit'):
+                await g.aclose()
+            DONE += 10
+
+        self.loop.run_until_complete(run())
+        self.assertEqual(DONE, 10)
+
     def test_async_gen_asyncio_asend_01(self):
         DONE = 0
 
@@ -801,6 +906,41 @@ class AsyncGenAsyncioTest(unittest.TestCase):
             self.loop.run_until_complete(run())
         self.assertEqual(DONE, 1)
 
+    def test_async_gen_asyncio_athrow_03(self):
+        DONE = 0
+
+        # test synchronous generators
+        def foo():
+            try:
+                yield
+            except:
+                pass
+        g = foo()
+        g.send(None)
+        with self.assertRaises(StopIteration):
+            g.throw(ValueError)
+
+        # now with asynchronous generators
+
+        async def gen():
+            nonlocal DONE
+            try:
+                yield
+            except:
+                pass
+            DONE = 1
+
+        async def run():
+            nonlocal DONE
+            g = gen()
+            await g.asend(None)
+            with self.assertRaises(StopAsyncIteration):
+                await g.athrow(ValueError)
+            DONE += 10
+
+        self.loop.run_until_complete(run())
+        self.assertEqual(DONE, 11)
+
     def test_async_gen_asyncio_athrow_tuple(self):
         async def gen():
             try:
diff --git a/Misc/NEWS b/Misc/NEWS
index 0b90319..68b9907 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -31,6 +31,9 @@ Core and Builtins
 
 - Issue #26182: Fix a refleak in code that raises DeprecationWarning.
 
+- Issue #28721: Fix asynchronous generators aclose() and athrow() to
+  handle StopAsyncIteration propagation properly.
+
 Library
 -------
 
diff --git a/Objects/genobject.c b/Objects/genobject.c
index ddf72cc..558f809 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -1931,9 +1931,17 @@ yield_close:
     return NULL;
 
 check_error:
-    if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)
-        || PyErr_ExceptionMatches(PyExc_GeneratorExit)
-    ) {
+    if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
+        o->agt_state = AWAITABLE_STATE_CLOSED;
+        if (o->agt_args == NULL) {
+            /* when aclose() is called we don't want to propagate
+               StopAsyncIteration; just raise StopIteration, signalling
+               that 'aclose()' is done. */
+            PyErr_Clear();
+            PyErr_SetNone(PyExc_StopIteration);
+        }
+    }
+    else if (PyErr_ExceptionMatches(PyExc_GeneratorExit)) {
         o->agt_state = AWAITABLE_STATE_CLOSED;
         PyErr_Clear();          /* ignore these errors */
         PyErr_SetNone(PyExc_StopIteration);
-- 
cgit v0.12