feat: add coroutine::await_in_coroutine to await awaitables in coroutine context

Author: wyfo

Cargo.toml

@@ -31,6 +31,9 @@ unindent = { version = "0.2.1", optional = true }
 # support crate for multiple-pymethods feature
 inventory = { version = "0.3.0", optional = true }
 
+# support crate for experimental async feature
+scoped-tls-hkt = { version = "0.1", optional = true }
+
 # crate integrations that can be added using the eponymous features
 anyhow = { version = "1.0.1", optional = true }
 chrono = { version = "0.4.25", default-features = false, optional = true }
@@ -41,7 +44,7 @@ hashbrown = { version = ">= 0.14.5, < 0.16", optional = true }
 indexmap = { version = ">= 2.5.0, < 3", optional = true }
 num-bigint = { version = "0.4.2", optional = true }
 num-complex = { version = ">= 0.4.6, < 0.5", optional = true }
-num-rational = {version = "0.4.1", optional = true }
+num-rational = { version = "0.4.1", optional = true }
 rust_decimal = { version = "1.15", default-features = false, optional = true }
 serde = { version = "1.0", optional = true }
 smallvec = { version = "1.0", optional = true }
@@ -63,7 +66,7 @@ rayon = "1.6.1"
 futures = "0.3.28"
 tempfile = "3.12.0"
 static_assertions = "1.1.0"
-uuid = {version = "1.10.0", features = ["v4"] }
+uuid = { version = "1.10.0", features = ["v4"] }
 
 [build-dependencies]
 pyo3-build-config = { path = "pyo3-build-config", version = "=0.23.3", features = ["resolve-config"] }
@@ -72,7 +75,7 @@ pyo3-build-config = { path = "pyo3-build-config", version = "=0.23.3", features
 default = ["macros"]
 
 # Enables support for `async fn` for `#[pyfunction]` and `#[pymethods]`.
-experimental-async = ["macros", "pyo3-macros/experimental-async"]
+experimental-async = ["macros", "pyo3-macros/experimental-async", "scoped-tls-hkt"]
 
 # Enables pyo3::inspect module and additional type information on FromPyObject
 # and IntoPy traits

guide/src/SUMMARY.md

@@ -25,6 +25,7 @@
   - [Mapping of Rust types to Python types](conversions/tables.md)
   - [Conversion traits](conversions/traits.md)
 - [Using `async` and `await`](async-await.md)
+  - [Awaiting Python awaitables](async-await/awaiting_python_awaitables)
 - [Parallelism](parallelism.md)
 - [Supporting Free-Threaded Python](free-threading.md)
 - [Debugging](debugging.md)

guide/src/async-await/awaiting_python_awaitables.md

@@ -0,0 +1,62 @@
+# Awaiting Python awaitables
+
+Python awaitable can be awaited on Rust side
+using [`await_in_coroutine`]({{#PYO3_DOCS_URL}}/pyo3/coroutine/function.await_in_coroutine).
+
+```rust
+# # ![allow(dead_code)]
+# #[cfg(feature = "experimental-async")] {
+use pyo3::{prelude::*, coroutine::await_in_coroutine};
+
+#[pyfunction]
+async fn wrap_awaitable(awaitable: PyObject) -> PyResult<PyObject> {
+    Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?.await
+}
+# }
+```
+
+Behind the scene, `await_in_coroutine` calls the `__await__` method of the Python awaitable (or `__iter__` for
+generator-based coroutine).
+
+## Restrictions
+
+As the name suggests, `await_in_coroutine` resulting future can only be awaited in coroutine context. Otherwise, it
+panics.
+
+```rust
+# # ![allow(dead_code)]
+# #[cfg(feature = "experimental-async")] {
+use pyo3::{prelude::*, coroutine::await_in_coroutine};
+
+#[pyfunction]
+fn block_on(awaitable: PyObject) -> PyResult<PyObject> {
+    let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?;
+    futures::executor::block_on(future) // ERROR: Python awaitable must be awaited in coroutine context
+}
+# }
+```
+
+The future must also be the only one to be awaited at a time; it means that it's forbidden to await it in a `select!`.
+Otherwise, it panics.
+
+```rust
+# # ![allow(dead_code)]
+# #[cfg(feature = "experimental-async")] {
+use futures::FutureExt;
+use pyo3::{prelude::*, coroutine::await_in_coroutine};
+
+#[pyfunction]
+async fn select(awaitable: PyObject) -> PyResult<PyObject> {
+    let future = Python::with_gil(|gil| await_in_coroutine(awaitable.bind(gil)))?;
+    futures::select_biased! {
+        _ = std::future::pending::<()>().fuse() => unreachable!(),
+        res = future.fuse() => res, // ERROR: Python awaitable mixed with Rust future
+    }
+}
+# }
+```
+
+These restrictions exist because awaiting a `await_in_coroutine` future strongly binds it to the
+enclosing coroutine. The coroutine will then delegate its `send`/`throw`/`close` methods to the
+awaited future. If it was awaited in a `select!`, `Coroutine::send` would no able to know if
+the value passed would have to be delegated or not.

newsfragments/3611.added.md

@@ -0,0 +1 @@
+Add `coroutine::await_in_coroutine` to await awaitables in coroutine context

pyo3-ffi/src/abstract_.rs

@@ -1,8 +1,9 @@
-use crate::object::*;
-use crate::pyport::Py_ssize_t;
+use std::os::raw::{c_char, c_int};
+
 #[cfg(any(Py_3_12, all(Py_3_8, not(Py_LIMITED_API))))]
 use libc::size_t;
-use std::os::raw::{c_char, c_int};
+
+use crate::{object::*, pyport::Py_ssize_t};
 
 #[inline]
 #[cfg(all(not(Py_3_13), not(PyPy)))] // CPython exposed as a function in 3.13, in object.h
@@ -143,7 +144,11 @@ extern "C" {
     pub fn PyIter_Next(arg1: *mut PyObject) -> *mut PyObject;
     #[cfg(all(not(PyPy), Py_3_10))]
     #[cfg_attr(PyPy, link_name = "PyPyIter_Send")]
-    pub fn PyIter_Send(iter: *mut PyObject, arg: *mut PyObject, presult: *mut *mut PyObject);
+    pub fn PyIter_Send(
+        iter: *mut PyObject,
+        arg: *mut PyObject,
+        presult: *mut *mut PyObject,
+    ) -> c_int;
 
     #[cfg_attr(PyPy, link_name = "PyPyNumber_Check")]
     pub fn PyNumber_Check(o: *mut PyObject) -> c_int;

src/coroutine.rs

@@ -11,28 +11,37 @@ use std::{
 use pyo3_macros::{pyclass, pymethods};
 
 use crate::{
-    coroutine::{cancel::ThrowCallback, waker::AsyncioWaker},
-    exceptions::{PyAttributeError, PyRuntimeError, PyStopIteration},
+    coroutine::waker::CoroutineWaker,
+    exceptions::{PyAttributeError, PyGeneratorExit, PyRuntimeError, PyStopIteration},
     panic::PanicException,
-    types::{string::PyStringMethods, PyIterator, PyString},
-    Bound, IntoPyObject, IntoPyObjectExt, Py, PyAny, PyErr, PyObject, PyResult, Python,
+    types::{string::PyStringMethods, PyString},
+    Bound, IntoPyObject, IntoPyObjectExt, Py, PyErr, PyObject, PyResult, Python,
 };
 
-pub(crate) mod cancel;
+mod asyncio;
+mod awaitable;
+mod cancel;
 mod waker;
 
-pub use cancel::CancelHandle;
+pub use awaitable::await_in_coroutine;
+pub use cancel::{CancelHandle, ThrowCallback};
 
 const COROUTINE_REUSED_ERROR: &str = "cannot reuse already awaited coroutine";
 
+pub(crate) enum CoroOp {
+    Send(PyObject),
+    Throw(PyObject),
+    Close,
+}
+
 /// Python coroutine wrapping a [`Future`].
 #[pyclass(crate = "crate")]
 pub struct Coroutine {
     name: Option<Py<PyString>>,
     qualname_prefix: Option<&'static str>,
     throw_callback: Option<ThrowCallback>,
     future: Option<Pin<Box<dyn Future<Output = PyResult<PyObject>> + Send>>>,
-    waker: Option<Arc<AsyncioWaker>>,
+    waker: Option<Arc<CoroutineWaker>>,
 }
 
 // Safety: `Coroutine` is allowed to be `Sync` even though the future is not,
@@ -71,55 +80,58 @@ impl Coroutine {
         }
     }
 
-    fn poll(&mut self, py: Python<'_>, throw: Option<PyObject>) -> PyResult<PyObject> {
+    fn poll_inner(&mut self, py: Python<'_>, mut op: CoroOp) -> PyResult<PyObject> {
         // raise if the coroutine has already been run to completion
         let future_rs = match self.future {
             Some(ref mut fut) => fut,
             None => return Err(PyRuntimeError::new_err(COROUTINE_REUSED_ERROR)),
         };
-        // reraise thrown exception it
-        match (throw, &self.throw_callback) {
-            (Some(exc), Some(cb)) => cb.throw(exc),
-            (Some(exc), None) => {
-                self.close();
-                return Err(PyErr::from_value(exc.into_bound(py)));
-            }
-            (None, _) => {}
+        // if the future is not pending on a Python awaitable,
+        // execute throw callback or complete on close
+        if !matches!(self.waker, Some(ref w) if w.is_delegated(py)) {
+            match op {
+                send @ CoroOp::Send(_) => op = send,
+                CoroOp::Throw(exc) => match &self.throw_callback {
+                    Some(cb) => {
+                        cb.throw(exc.clone_ref(py));
+                        op = CoroOp::Send(py.None());
+                    }
+                    None => return Err(PyErr::from_value(exc.into_bound(py))),
+                },
+                CoroOp::Close => return Err(PyGeneratorExit::new_err(py.None())),
+            };
         }
         // create a new waker, or try to reset it in place
         if let Some(waker) = self.waker.as_mut().and_then(Arc::get_mut) {
-            waker.reset();
+            waker.reset(op);
         } else {
-            self.waker = Some(Arc::new(AsyncioWaker::new()));
+            self.waker = Some(Arc::new(CoroutineWaker::new(op)));
         }
-        let waker = Waker::from(self.waker.clone().unwrap());
-        // poll the Rust future and forward its results if ready
+        // poll the future and forward its results if ready; otherwise, yield from waker
         // polling is UnwindSafe because the future is dropped in case of panic
+        let waker = Waker::from(self.waker.clone().unwrap());
         let poll = || future_rs.as_mut().poll(&mut Context::from_waker(&waker));
         match panic::catch_unwind(panic::AssertUnwindSafe(poll)) {
-            Ok(Poll::Ready(res)) => {
-                self.close();
-                return Err(PyStopIteration::new_err((res?,)));
-            }
-            Err(err) => {
-                self.close();
-                return Err(PanicException::from_panic_payload(err));
-            }
-            _ => {}
+            Err(err) => Err(PanicException::from_panic_payload(err)),
+            // See #4407, `PyStopIteration::new_err` argument must be wrap in a tuple,
+            // otherwise, when a tuple is returned, its fields would be expanded as error
+            // arguments
+            Ok(Poll::Ready(res)) => Err(PyStopIteration::new_err((res?,))),
+            Ok(Poll::Pending) => match self.waker.as_ref().unwrap().yield_(py) {
+                Ok(to_yield) => Ok(to_yield),
+                Err(err) => Err(err),
+            },
         }
-        // otherwise, initialize the waker `asyncio.Future`
-        if let Some(future) = self.waker.as_ref().unwrap().initialize_future(py)? {
-            // `asyncio.Future` must be awaited; fortunately, it implements `__iter__ = __await__`
-            // and will yield itself if its result has not been set in polling above
-            if let Some(future) = PyIterator::from_object(future).unwrap().next() {
-                // future has not been leaked into Python for now, and Rust code can only call
-                // `set_result(None)` in `Wake` implementation, so it's safe to unwrap
-                return Ok(future.unwrap().into());
-            }
+    }
+
+    fn poll(&mut self, py: Python<'_>, op: CoroOp) -> PyResult<PyObject> {
+        let result = self.poll_inner(py, op);
+        if result.is_err() {
+            // the Rust future is dropped, and the field set to `None`
+            // to indicate the coroutine has been run to completion
+            drop(self.future.take());
         }
-        // if waker has been waken during future polling, this is roughly equivalent to
-        // `await asyncio.sleep(0)`, so just yield `None`.
-        Ok(py.None())
+        result
     }
 }
 
@@ -145,25 +157,27 @@ impl Coroutine {
         }
     }
 
-    fn send(&mut self, py: Python<'_>, _value: &Bound<'_, PyAny>) -> PyResult<PyObject> {
-        self.poll(py, None)
+    fn send(&mut self, py: Python<'_>, value: PyObject) -> PyResult<PyObject> {
+        self.poll(py, CoroOp::Send(value))
     }
 
     fn throw(&mut self, py: Python<'_>, exc: PyObject) -> PyResult<PyObject> {
-        self.poll(py, Some(exc))
+        self.poll(py, CoroOp::Throw(exc))
     }
 
-    fn close(&mut self) {
-        // the Rust future is dropped, and the field set to `None`
-        // to indicate the coroutine has been run to completion
-        drop(self.future.take());
+    fn close(&mut self, py: Python<'_>) -> PyResult<()> {
+        match self.poll(py, CoroOp::Close) {
+            Ok(_) => Ok(()),
+            Err(err) if err.is_instance_of::<PyGeneratorExit>(py) => Ok(()),
+            Err(err) => Err(err),
+        }
     }
 
     fn __await__(self_: Py<Self>) -> Py<Self> {
         self_
     }
 
     fn __next__(&mut self, py: Python<'_>) -> PyResult<PyObject> {
-        self.poll(py, None)
+        self.poll(py, CoroOp::Send(py.None()))
     }
 }