Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion Lib/asyncio/futures.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
4 changes: 4 additions & 0 deletions Lib/test/test_asyncio/test_futures.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Expose the future being waited on by a ``FutureIter`` as ``fi_future``.
17 changes: 17 additions & 0 deletions Modules/_asynciomodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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},
Expand Down
Loading