gh-158192: expose the Future in FutureIter as fi_future

I'd be happy to bikeshed about the name if anybody wants to, but the
scheme kind of matches generators and similar.

For consistency, I made the Python implementation match, which
required a wrapper class. I can drop that if it doesn't seem worth the
hassle; the future can be introspected from the frame object of the
generator already.

Author: msullivan

Lib/asyncio/futures.py

@@ -293,17 +293,43 @@ def set_exception(self, exception):
         self.__schedule_callbacks()
         self.__log_traceback = True
 
-    def __await__(self):
+    def _await(self):
         if not self.done():
             self._asyncio_future_blocking = True
             yield self  # This tells Task to wait for completion.
         if not self.done():
             raise RuntimeError("await wasn't used with future")
         return self.result()  # May raise too.
 
+    def __await__(self):
+        return _FutureIter(self)
+
     __iter__ = __await__  # make compatible with 'yield from'.
 
 
+class _FutureIter:
+    def __init__(self, fut):
+        # The whole point of this class is to expose fi_future, like the C
+        # version
+        self.fi_future = fut
+        self.__gen = fut._await()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return self.__gen.__next__()
+
+    def send(self, v, /):
+        return self.__gen.send(v)
+
+    def throw(self, *args):
+        return self.__gen.throw(*args)
+
+    def close(self):
+        self.__gen.close()
+
+
 # Needed for testing purposes.
 _PyFuture = Future
 

Lib/test/test_asyncio/test_futures.py

@@ -712,6 +712,10 @@ def test_future_stop_iteration_args(self):
             self.fail('StopIteration was expected')
         self.assertEqual(result, (1, 2))
 
+    def test_future_iter_fut(self):
+        fut = self._new_future(loop=self.loop)
+        self.assertIs(iter(fut).fi_future, fut)
+
     def test_future_iter_throw(self):
         fut = self._new_future(loop=self.loop)
         fi = iter(fut)

Modules/_asynciomodule.c

@@ -1971,6 +1971,22 @@ static PyMethodDef FutureIter_methods[] = {
     {NULL, NULL}        /* Sentinel */
 };
 
+static PyObject *
+FutureIter_get_fi_future(PyObject *op, void *Py_UNUSED(closure))
+{
+    futureiterobject *it = (futureiterobject*)op;
+    PyObject *fut;
+    Py_BEGIN_CRITICAL_SECTION(op);
+    fut = it->future ? Py_NewRef(it->future) : Py_None;
+    Py_END_CRITICAL_SECTION();
+    return fut;
+}
+
+static PyGetSetDef FutureIter_getsetlist[] = {
+    {"fi_future", FutureIter_get_fi_future, NULL, NULL},
+    {NULL} /* Sentinel */
+};
+
 static PyType_Slot FutureIter_slots[] = {
     {Py_tp_dealloc, FutureIter_dealloc},
     {Py_tp_getattro, PyObject_GenericGetAttr},
@@ -1979,6 +1995,7 @@ static PyType_Slot FutureIter_slots[] = {
     {Py_tp_iter, PyObject_SelfIter},
     {Py_tp_iternext, FutureIter_iternext},
     {Py_tp_methods, FutureIter_methods},
+    {Py_tp_getset, FutureIter_getsetlist},
 
     // async methods
     {Py_am_send, FutureIter_am_send},