Issue #24450: Add gi_yieldfrom to generators; cr_await to coroutines.

Patch by Benno Leslie and Yury Selivanov.

2 files changed, 66 insertions, 1 deletions

diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
index 3413a12..9d97123 100644
--- a/Lib/test/test_coroutines.py
+++ b/Lib/test/test_coroutines.py
@@ -350,6 +350,36 @@ class CoroutineTest(unittest.TestCase):
                                     "coroutine ignored GeneratorExit"):
             c.close()
 
+    def test_cr_await(self):
+        @types.coroutine
+        def a():
+            self.assertEqual(inspect.getcoroutinestate(coro_b), inspect.CORO_RUNNING)
+            self.assertIsNone(coro_b.cr_await)
+            yield
+            self.assertEqual(inspect.getcoroutinestate(coro_b), inspect.CORO_RUNNING)
+            self.assertIsNone(coro_b.cr_await)
+
+        async def c():
+            await a()
+
+        async def b():
+            self.assertIsNone(coro_b.cr_await)
+            await c()
+            self.assertIsNone(coro_b.cr_await)
+
+        coro_b = b()
+        self.assertEqual(inspect.getcoroutinestate(coro_b), inspect.CORO_CREATED)
+        self.assertIsNone(coro_b.cr_await)
+
+        coro_b.send(None)
+        self.assertEqual(inspect.getcoroutinestate(coro_b), inspect.CORO_SUSPENDED)
+        self.assertEqual(coro_b.cr_await.cr_await.gi_code.co_name, 'a')
+
+        with self.assertRaises(StopIteration):
+            coro_b.send(None)  # complete coroutine
+        self.assertEqual(inspect.getcoroutinestate(coro_b), inspect.CORO_CLOSED)
+        self.assertIsNone(coro_b.cr_await)
+
     def test_corotype_1(self):
         ct = types.CoroutineType
         self.assertIn('into coroutine', ct.send.__doc__)
diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py
index fe4b138..25cc628 100644
--- a/Lib/test/test_generators.py
+++ b/Lib/test/test_generators.py
@@ -3,6 +3,8 @@ import sys
 import unittest
 import warnings
 import weakref
+import inspect
+import types
 
 from test import support
 
@@ -259,6 +261,39 @@ class ExceptionTest(unittest.TestCase):
             next(g)
 
 
+class YieldFromTests(unittest.TestCase):
+    def test_generator_gi_yieldfrom(self):
+        def a():
+            self.assertEqual(inspect.getgeneratorstate(gen_b), inspect.GEN_RUNNING)
+            self.assertIsNone(gen_b.gi_yieldfrom)
+            yield
+            self.assertEqual(inspect.getgeneratorstate(gen_b), inspect.GEN_RUNNING)
+            self.assertIsNone(gen_b.gi_yieldfrom)
+
+        def b():
+            self.assertIsNone(gen_b.gi_yieldfrom)
+            yield from a()
+            self.assertIsNone(gen_b.gi_yieldfrom)
+            yield
+            self.assertIsNone(gen_b.gi_yieldfrom)
+
+        gen_b = b()
+        self.assertEqual(inspect.getgeneratorstate(gen_b), inspect.GEN_CREATED)
+        self.assertIsNone(gen_b.gi_yieldfrom)
+
+        gen_b.send(None)
+        self.assertEqual(inspect.getgeneratorstate(gen_b), inspect.GEN_SUSPENDED)
+        self.assertEqual(gen_b.gi_yieldfrom.gi_code.co_name, 'a')
+
+        gen_b.send(None)
+        self.assertEqual(inspect.getgeneratorstate(gen_b), inspect.GEN_SUSPENDED)
+        self.assertIsNone(gen_b.gi_yieldfrom)
+
+        [] = gen_b  # Exhaust generator
+        self.assertEqual(inspect.getgeneratorstate(gen_b), inspect.GEN_CLOSED)
+        self.assertIsNone(gen_b.gi_yieldfrom)
+
+
 tutorial_tests = """
 Let's try a simple generator:
 
@@ -624,7 +659,7 @@ From the Iterators list, about the types of these things.
 >>> type(i)
 <class 'generator'>
 >>> [s for s in dir(i) if not s.startswith('_')]
-['close', 'gi_code', 'gi_frame', 'gi_running', 'send', 'throw']
+['close', 'gi_code', 'gi_frame', 'gi_running', 'gi_yieldfrom', 'send', 'throw']
 >>> from test.support import HAVE_DOCSTRINGS
 >>> print(i.__next__.__doc__ if HAVE_DOCSTRINGS else 'Implement next(self).')
 Implement next(self).