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

1st1 · 1st1 · commit e13f8f3cabda · 2015-07-03T00:23:30.000-04:00
Patch by Benno Leslie and Yury Selivanov.
diff --git a/Doc/library/inspect.rst b/Doc/library/inspect.rst
@@ -182,12 +182,19 @@ attributes:
 +-----------+-----------------+---------------------------+
 |           | __qualname__    | qualified name            |
 +-----------+-----------------+---------------------------+
+|           | cr_await        | object being awaited on,  |
+|           |                 | or ``None``               |
++-----------+-----------------+---------------------------+
 |           | cr_frame        | frame                     |
 +-----------+-----------------+---------------------------+
 |           | cr_running      | is the coroutine running? |
 +-----------+-----------------+---------------------------+
 |           | cr_code         | code                      |
 +-----------+-----------------+---------------------------+
+|           | gi_yieldfrom    | object being iterated by  |
+|           |                 | ``yield from``, or        |
+|           |                 | ``None``                  |
++-----------+-----------------+---------------------------+
 | builtin   | __doc__         | documentation string      |
 +-----------+-----------------+---------------------------+
 |           | __name__        | original name of this     |
diff --git a/Doc/whatsnew/3.5.rst b/Doc/whatsnew/3.5.rst
@@ -84,6 +84,9 @@ New built-in features:
 * ``b'\xf0\x9f\x90\x8d'.hex()``, ``bytearray(b'\xf0\x9f\x90\x8d').hex()``,
   ``memoryview(b'\xf0\x9f\x90\x8d').hex()``: :issue:`9951` - A ``hex`` method
   has been added to bytes, bytearray, and memoryview.
+* Generators have new ``gi_yieldfrom`` attribute, which returns the
+  object being iterated by ``yield from`` expressions. (Contributed
+  by Benno Leslie and Yury Selivanov in :issue:`24450`.)
 
 Implementation improvements:
 
diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
@@ -350,6 +350,36 @@ async def coro():
                                     "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
@@ -3,6 +3,8 @@
 import unittest
 import warnings
 import weakref
+import inspect
+import types
 
 from test import support
 
@@ -259,6 +261,39 @@ def f():
             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 @@ def f():
 >>> 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).
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -27,6 +27,9 @@ Core and Builtins
   used in types.coroutine to be instance of collections.abc.Generator;
   inspect.isawaitable was removed (use collections.abc.Awaitable).
 
+- Issue #24450: Add gi_yieldfrom to generators and cr_await to coroutines.
+  Contributed by Benno Leslie and Yury Selivanov.
+
 Library
 -------
 
diff --git a/Objects/genobject.c b/Objects/genobject.c
@@ -552,11 +552,22 @@ gen_set_qualname(PyGenObject *op, PyObject *value)
     return 0;
 }
 
+static PyObject *
+gen_getyieldfrom(PyGenObject *gen)
+{
+    PyObject *yf = gen_yf(gen);
+    if (yf == NULL)
+        Py_RETURN_NONE;
+    return yf;
+}
+
 static PyGetSetDef gen_getsetlist[] = {
     {"__name__", (getter)gen_get_name, (setter)gen_set_name,
      PyDoc_STR("name of the generator")},
     {"__qualname__", (getter)gen_get_qualname, (setter)gen_set_qualname,
      PyDoc_STR("qualified name of the generator")},
+    {"gi_yieldfrom", (getter)gen_getyieldfrom, NULL,
+     PyDoc_STR("object being iterated by yield from, or None")},
     {NULL} /* Sentinel */
 };
 
@@ -776,11 +787,22 @@ coro_await(PyCoroObject *coro)
     return (PyObject *)cw;
 }
 
+static PyObject *
+coro_get_cr_await(PyCoroObject *coro)
+{
+    PyObject *yf = gen_yf((PyGenObject *) coro);
+    if (yf == NULL)
+        Py_RETURN_NONE;
+    return yf;
+}
+
 static PyGetSetDef coro_getsetlist[] = {
     {"__name__", (getter)gen_get_name, (setter)gen_set_name,
      PyDoc_STR("name of the coroutine")},
     {"__qualname__", (getter)gen_get_qualname, (setter)gen_set_qualname,
      PyDoc_STR("qualified name of the coroutine")},
+    {"cr_await", (getter)coro_get_cr_await, NULL,
+     PyDoc_STR("object being awaited on, or None")},
     {NULL} /* Sentinel */
 };
 
