From 55a0e82b7b4797b2298cebde85ad83064a63f180 Mon Sep 17 00:00:00 2001 From: Joseph Perez Date: Thu, 30 Nov 2023 10:12:49 +0100 Subject: [PATCH] feat: add `PyFuture` to await Python awaitables --- guide/src/SUMMARY.md | 1 + guide/src/async-await/pyfuture.md | 51 ++++++++ pyo3-ffi/src/abstract_.rs | 6 +- src/coroutine.rs | 97 +++++++-------- src/coroutine/asyncio.rs | 90 ++++++++++++++ src/coroutine/waker.rs | 157 ++++++++++++------------ src/lib.rs | 1 + src/types/future.rs | 195 ++++++++++++++++++++++++++++++ src/types/mod.rs | 2 + tests/test_awaitable.rs | 134 ++++++++++++++++++++ tests/test_coroutine.rs | 2 +- 11 files changed, 605 insertions(+), 131 deletions(-) create mode 100644 guide/src/async-await/pyfuture.md create mode 100644 src/coroutine/asyncio.rs create mode 100644 src/types/future.rs create mode 100644 tests/test_awaitable.rs diff --git a/guide/src/SUMMARY.md b/guide/src/SUMMARY.md index 9e738a79946..00976562c60 100644 --- a/guide/src/SUMMARY.md +++ b/guide/src/SUMMARY.md @@ -20,6 +20,7 @@ - [Python exceptions](exception.md) - [Calling Python from Rust](python_from_rust.md) - [Using `async` and `await`](async-await.md) + - [Awaiting Python awaitables](async-await/pyfuture.md) - [GIL, mutability and object types](types.md) - [Parallelism](parallelism.md) - [Debugging](debugging.md) diff --git a/guide/src/async-await/pyfuture.md b/guide/src/async-await/pyfuture.md new file mode 100644 index 00000000000..8a82aaee0ac --- /dev/null +++ b/guide/src/async-await/pyfuture.md @@ -0,0 +1,51 @@ +# Awaiting Python awaitables + +Python awaitable can be awaited on Rust side using [`PyFuture`]({{#PYO3_DOCS_URL}}/pyo3/types/struct.PyFuture.html). + +```rust +# #![allow(dead_code)] +use pyo3::{prelude::*, types::PyFuture}; + +#[pyfunction] +async fn wrap_awaitable(awaitable: PyObject) -> PyResult { + let future: Py = Python::with_gil(|gil| Py::from_object(gil, awaitable))?; + future.await +} +``` + +`PyFuture::from_object` construct a `PyFuture` from a Python awaitable object, by calling its `__await__` method (or `__iter__` for generator-based coroutine). + +## Restrictions + +`PyFuture` can only be awaited in the context of a PyO3 coroutine. Otherwise, it panics. + +```rust +# #![allow(dead_code)] +use pyo3::{prelude::*, types::PyFuture}; + +#[pyfunction] +fn block_on(awaitable: PyObject) -> PyResult { + let future: Py = Python::with_gil(|gil| Py::from_object(gil, awaitable))?; + futures::executor::block_on(future) // ERROR: PyFuture must be awaited in coroutine context +} +``` + +`PyFuture` must be the only Rust future awaited; it means that it's forbidden to `select!` a `Pyfuture`. Otherwise, it panics. + +```rust +# #![allow(dead_code)] +use std::future; +use futures::FutureExt; +use pyo3::{prelude::*, types::PyFuture}; + +#[pyfunction] +async fn select(awaitable: PyObject) -> PyResult { + let future: Py = Python::with_gil(|gil| Py::from_object(gil, awaitable))?; + futures::select_biased! { + _ = future::pending::<()>().fuse() => unreachable!(), + res = future.fuse() => res, // ERROR: Python awaitable mixed with Rust future + } +} +``` + +These restrictions exist because awaiting a `PyFuture` strongly binds it to the enclosing coroutine. The coroutine will then delegate its `send`/`throw`/`close` methods to the awaited `PyFuture`. If it was awaited in a `select!`, `Coroutine::send` would no able to know if the value passed would have to be delegated to the `Pyfuture` or not. diff --git a/pyo3-ffi/src/abstract_.rs b/pyo3-ffi/src/abstract_.rs index 0b3b7dbb3c2..6e4f146aed4 100644 --- a/pyo3-ffi/src/abstract_.rs +++ b/pyo3-ffi/src/abstract_.rs @@ -128,7 +128,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; diff --git a/src/coroutine.rs b/src/coroutine.rs index 1d7d9f67d7a..1e7c2922209 100644 --- a/src/coroutine.rs +++ b/src/coroutine.rs @@ -1,26 +1,26 @@ //! Python coroutine implementation, used notably when wrapping `async fn` //! with `#[pyfunction]`/`#[pymethods]`. -use std::task::Waker; use std::{ future::Future, panic, pin::Pin, sync::Arc, - task::{Context, Poll}, + task::{Context, Poll, Waker}, }; use pyo3_macros::{pyclass, pymethods}; use crate::{ - coroutine::waker::AsyncioWaker, + coroutine::waker::CoroutineWaker, exceptions::{PyAttributeError, PyRuntimeError, PyStopIteration}, pyclass::IterNextOutput, - types::{PyIterator, PyString}, - IntoPy, Py, PyAny, PyErr, PyObject, PyResult, Python, + types::PyString, + IntoPy, Py, PyErr, PyObject, PyResult, Python, }; +mod asyncio; pub(crate) mod cancel; -mod waker; +pub(crate) mod waker; use crate::coroutine::cancel::ThrowCallback; use crate::panic::PanicException; @@ -36,7 +36,7 @@ pub struct Coroutine { throw_callback: Option, allow_threads: bool, future: Option> + Send>>>, - waker: Option>, + waker: Option>, } impl Coroutine { @@ -73,33 +73,37 @@ impl Coroutine { } } - fn poll( + fn poll_inner( &mut self, py: Python<'_>, - throw: Option, + mut sent_result: Option>, ) -> PyResult> { // 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.as_ref(py)), - (Some(exc), None) => { - self.close(); - return Err(PyErr::from_value(exc.as_ref(py))); + // 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.yielded_from_awaitable(py)) { + match (sent_result, &self.throw_callback) { + (res @ Some(Ok(_)), _) => sent_result = res, + (Some(Err(err)), Some(cb)) => { + cb.throw(err.as_ref(py)); + sent_result = Some(Ok(py.None().into())); + } + (Some(Err(err)), None) => return Err(PyErr::from_value(err.as_ref(py))), + (None, _) => return Ok(IterNextOutput::Return(py.None().into())), } - _ => {} } // 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(sent_result); } else { - self.waker = Some(Arc::new(AsyncioWaker::new())); + self.waker = Some(Arc::new(CoroutineWaker::new(sent_result))); } let waker = Waker::from(self.waker.clone().unwrap()); - // poll the Rust future and forward its results if ready + // poll the Rust future and forward its results if ready; otherwise, yield from waker // polling is UnwindSafe because the future is dropped in case of panic let poll = || { if self.allow_threads { @@ -109,29 +113,27 @@ impl Coroutine { } }; match panic::catch_unwind(panic::AssertUnwindSafe(poll)) { - Ok(Poll::Ready(res)) => { - self.close(); - return Ok(IterNextOutput::Return(res?)); - } - Err(err) => { - self.close(); - return Err(PanicException::from_panic_payload(err)); - } - _ => {} + Err(err) => Err(PanicException::from_panic_payload(err)), + Ok(Poll::Ready(res)) => Ok(IterNextOutput::Return(res?)), + Ok(Poll::Pending) => match self.waker.as_ref().unwrap().yield_(py) { + Ok(to_yield) => Ok(IterNextOutput::Yield(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(IterNextOutput::Yield(future.unwrap().into())); - } + } + + fn poll( + &mut self, + py: Python<'_>, + sent_result: Option>, + ) -> PyResult> { + let result = self.poll_inner(py, sent_result); + if matches!(result, Ok(IterNextOutput::Return(_)) | 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(IterNextOutput::Yield(py.None().into())) + result } } @@ -163,18 +165,17 @@ impl Coroutine { } } - fn send(&mut self, py: Python<'_>, _value: &PyAny) -> PyResult { - iter_result(self.poll(py, None)?) + fn send(&mut self, py: Python<'_>, value: PyObject) -> PyResult { + iter_result(self.poll(py, Some(Ok(value)))?) } fn throw(&mut self, py: Python<'_>, exc: PyObject) -> PyResult { - iter_result(self.poll(py, Some(exc))?) + iter_result(self.poll(py, Some(Err(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<()> { + self.poll(py, None)?; + Ok(()) } fn __await__(self_: Py) -> Py { @@ -182,6 +183,6 @@ impl Coroutine { } fn __next__(&mut self, py: Python<'_>) -> PyResult> { - self.poll(py, None) + self.poll(py, Some(Ok(py.None().into()))) } } diff --git a/src/coroutine/asyncio.rs b/src/coroutine/asyncio.rs new file mode 100644 index 00000000000..4581fed28b9 --- /dev/null +++ b/src/coroutine/asyncio.rs @@ -0,0 +1,90 @@ +//! Coroutine implementation compatible with asyncio. +use crate::sync::GILOnceCell; +use crate::types::{PyCFunction, PyIterator}; +use crate::{intern, wrap_pyfunction, IntoPy, Py, PyAny, PyObject, PyResult, Python}; +use pyo3_macros::pyfunction; + +/// `asyncio.get_running_loop` +fn get_running_loop(py: Python<'_>) -> PyResult<&PyAny> { + static GET_RUNNING_LOOP: GILOnceCell = GILOnceCell::new(); + let import = || -> PyResult<_> { + let module = py.import("asyncio")?; + Ok(module.getattr("get_running_loop")?.into()) + }; + GET_RUNNING_LOOP + .get_or_try_init(py, import)? + .as_ref(py) + .call0() +} + +/// Asyncio-compatible coroutine waker. +/// +/// Polling a Rust future yields an `asyncio.Future`, whose `set_result` method is called +/// when `Waker::wake` is called. +pub(super) struct AsyncioWaker { + event_loop: PyObject, + future: PyObject, +} + +impl AsyncioWaker { + pub(super) fn new(py: Python<'_>) -> PyResult { + let event_loop = get_running_loop(py)?.into_py(py); + let future = event_loop.call_method0(py, "create_future")?; + Ok(Self { event_loop, future }) + } + + pub(super) fn yield_(&self, py: Python<'_>) -> PyResult { + let __await__; + // `asyncio.Future` must be awaited; in normal case, it implements `__iter__ = __await__`, + // but `create_future` may have been overriden + let mut iter = match PyIterator::from_object(self.future.as_ref(py)) { + Ok(iter) => iter, + Err(_) => { + __await__ = self.future.call_method0(py, intern!(py, "__await__"))?; + PyIterator::from_object(__await__.as_ref(py))? + } + }; + // future has not been waken (because `yield_waken` would have been called + // otherwise), so it is expected to yield itself + Ok(iter.next().expect("future didn't yield")?.into_py(py)) + } + + pub(super) fn yield_waken(py: Python<'_>) -> PyResult { + Ok(py.None().into()) + } + + pub(super) fn wake(&self, py: Python<'_>) -> PyResult<()> { + static RELEASE_WAITER: GILOnceCell> = GILOnceCell::new(); + let release_waiter = RELEASE_WAITER + .get_or_try_init(py, || wrap_pyfunction!(release_waiter, py).map(Into::into))?; + // `Future.set_result` must be called in event loop thread, + // so it requires `call_soon_threadsafe` + let call_soon_threadsafe = self.event_loop.call_method1( + py, + intern!(py, "call_soon_threadsafe"), + (release_waiter, self.future.as_ref(py)), + ); + if let Err(err) = call_soon_threadsafe { + // `call_soon_threadsafe` will raise if the event loop is closed; + // instead of catching an unspecific `RuntimeError`, check directly if it's closed. + let is_closed = self.event_loop.call_method0(py, "is_closed")?; + if !is_closed.extract(py)? { + return Err(err); + } + } + Ok(()) + } +} + +/// Call `future.set_result` if the future is not done. +/// +/// Future can be cancelled by the event loop before being waken. +/// See +#[pyfunction(crate = "crate")] +fn release_waiter(future: &PyAny) -> PyResult<()> { + let done = future.call_method0(intern!(future.py(), "done"))?; + if !done.extract::()? { + future.call_method1(intern!(future.py(), "set_result"), (future.py().None(),))?; + } + Ok(()) +} diff --git a/src/coroutine/waker.rs b/src/coroutine/waker.rs index 8a1166ce3fb..0fb8c56eae9 100644 --- a/src/coroutine/waker.rs +++ b/src/coroutine/waker.rs @@ -1,101 +1,96 @@ +use crate::coroutine::asyncio::AsyncioWaker; +use crate::exceptions::PyStopIteration; +use crate::pyclass::IterNextOutput; use crate::sync::GILOnceCell; -use crate::types::PyCFunction; -use crate::{intern, wrap_pyfunction, Py, PyAny, PyObject, PyResult, Python}; -use pyo3_macros::pyfunction; +use crate::types::PyFuture; +use crate::{intern, Py, PyObject, PyResult, Python}; +use std::cell::Cell; use std::sync::Arc; -use std::task::Wake; +use std::task::{Poll, Wake}; -/// Lazy `asyncio.Future` wrapper, implementing [`Wake`] by calling `Future.set_result`. -/// -/// asyncio future is let uninitialized until [`initialize_future`][1] is called. -/// If [`wake`][2] is called before future initialization (during Rust future polling), -/// [`initialize_future`][1] will return `None` (it is roughly equivalent to `asyncio.sleep(0)`) -/// -/// [1]: AsyncioWaker::initialize_future -/// [2]: AsyncioWaker::wake -pub struct AsyncioWaker(GILOnceCell>); +const MIXED_AWAITABLE_AND_FUTURE_ERROR: &str = "Python awaitable mixed with Rust future"; -impl AsyncioWaker { - pub(super) fn new() -> Self { - Self(GILOnceCell::new()) - } +pub(crate) enum FutureOrPoll { + Future(Py), + Poll(Poll>), +} - pub(super) fn reset(&mut self) { - self.0.take(); - } +thread_local! { + pub(crate) static FUTURE_OR_POLL: Cell> = Cell::new(None); +} - pub(super) fn initialize_future<'a>(&'a self, py: Python<'a>) -> PyResult> { - let init = || LoopAndFuture::new(py).map(Some); - let loop_and_future = self.0.get_or_try_init(py, init)?.as_ref(); - Ok(loop_and_future.map(|LoopAndFuture { future, .. }| future.as_ref(py))) - } +enum State { + Waken, + Yielded(PyObject), + Pending(AsyncioWaker), } -impl Wake for AsyncioWaker { - fn wake(self: Arc) { - self.wake_by_ref() - } +pub(super) struct CoroutineWaker { + state: GILOnceCell, + sent_result: Option>, +} - fn wake_by_ref(self: &Arc) { - Python::with_gil(|gil| { - if let Some(loop_and_future) = self.0.get_or_init(gil, || None) { - loop_and_future - .set_result(gil) - .expect("unexpected error in coroutine waker"); - } - }); +impl CoroutineWaker { + pub(super) fn new(sent_result: Option>) -> Self { + Self { + state: GILOnceCell::new(), + sent_result, + } } -} -struct LoopAndFuture { - event_loop: PyObject, - future: PyObject, -} + pub(super) fn reset(&mut self, sent_result: Option>) { + self.state.take(); + self.sent_result = sent_result; + } -impl LoopAndFuture { - fn new(py: Python<'_>) -> PyResult { - static GET_RUNNING_LOOP: GILOnceCell = GILOnceCell::new(); - let import = || -> PyResult<_> { - let module = py.import("asyncio")?; - Ok(module.getattr("get_running_loop")?.into()) - }; - let event_loop = GET_RUNNING_LOOP.get_or_try_init(py, import)?.call0(py)?; - let future = event_loop.call_method0(py, "create_future")?; - Ok(Self { event_loop, future }) + pub(super) fn yielded_from_awaitable(&self, py: Python<'_>) -> bool { + matches!(self.state.get(py), Some(State::Yielded(_))) } - fn set_result(&self, py: Python<'_>) -> PyResult<()> { - static RELEASE_WAITER: GILOnceCell> = GILOnceCell::new(); - let release_waiter = RELEASE_WAITER - .get_or_try_init(py, || wrap_pyfunction!(release_waiter, py).map(Into::into))?; - // `Future.set_result` must be called in event loop thread, - // so it requires `call_soon_threadsafe` - let call_soon_threadsafe = self.event_loop.call_method1( - py, - intern!(py, "call_soon_threadsafe"), - (release_waiter, self.future.as_ref(py)), - ); - if let Err(err) = call_soon_threadsafe { - // `call_soon_threadsafe` will raise if the event loop is closed; - // instead of catching an unspecific `RuntimeError`, check directly if it's closed. - let is_closed = self.event_loop.call_method0(py, "is_closed")?; - if !is_closed.extract(py)? { - return Err(err); - } + pub(super) fn yield_(&self, py: Python<'_>) -> PyResult { + let init = || PyResult::Ok(State::Pending(AsyncioWaker::new(py)?)); + let state = self.state.get_or_try_init(py, init)?; + match state { + State::Waken => AsyncioWaker::yield_waken(py), + State::Yielded(obj) => Ok(obj.clone()), + State::Pending(waker) => waker.yield_(py), } - Ok(()) } } -/// Call `future.set_result` if the future is not done. -/// -/// Future can be cancelled by the event loop before being waken. -/// See -#[pyfunction(crate = "crate")] -fn release_waiter(future: &PyAny) -> PyResult<()> { - let done = future.call_method0(intern!(future.py(), "done"))?; - if !done.extract::()? { - future.call_method1(intern!(future.py(), "set_result"), (future.py().None(),))?; +impl Wake for CoroutineWaker { + fn wake(self: Arc) { + self.wake_by_ref() + } + + fn wake_by_ref(self: &Arc) { + Python::with_gil(|gil| match FUTURE_OR_POLL.take() { + Some(FutureOrPoll::Future(fut)) => FUTURE_OR_POLL.set(Some(FutureOrPoll::Poll( + match fut.as_ref(gil).next(&self.sent_result) { + Ok(IterNextOutput::Return(ret)) => Poll::Ready(Ok(ret.into())), + Ok(IterNextOutput::Yield(yielded)) => { + if self.state.set(gil, State::Yielded(yielded.into())).is_err() { + panic!("{}", MIXED_AWAITABLE_AND_FUTURE_ERROR) + } + Poll::Pending + } + Err(err) if err.is_instance_of::(gil) => Poll::Ready( + err.value(gil) + .getattr(intern!(gil, "value")) + .map(Into::into), + ), + Err(err) => Poll::Ready(Err(err)), + }, + ))), + Some(FutureOrPoll::Poll(_)) => unreachable!(), + None => match self.state.get_or_init(gil, || State::Waken) { + State::Waken => {} + State::Yielded(_) => { + println!("{}", std::backtrace::Backtrace::force_capture()); + panic!("{}", MIXED_AWAITABLE_AND_FUTURE_ERROR) + } + State::Pending(waker) => waker.wake(gil).expect("wake error"), + }, + }) } - Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index 762644f1a2c..b64caef8858 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -467,6 +467,7 @@ pub mod doc_test { "README.md" => readme_md, "guide/src/advanced.md" => guide_advanced_md, "guide/src/async-await.md" => guide_async_await_md, + "guide/src/async-await/pyfuture.md" => guide_async_await_pyfuture_md, "guide/src/building_and_distribution.md" => guide_building_and_distribution_md, "guide/src/building_and_distribution/multiple_python_versions.md" => guide_bnd_multiple_python_versions_md, "guide/src/class.md" => guide_class_md, diff --git a/src/types/future.rs b/src/types/future.rs new file mode 100644 index 00000000000..3b5772ccf12 --- /dev/null +++ b/src/types/future.rs @@ -0,0 +1,195 @@ +use crate::coroutine::waker::{FutureOrPoll, FUTURE_OR_POLL}; +use crate::exceptions::{PyAttributeError, PyTypeError}; +use crate::instance::Py2; +use crate::pyclass::IterNextOutput; +use crate::sync::GILOnceCell; +use crate::types::any::PyAnyMethods; +use crate::{ + ffi, IntoPy, Py, PyAny, PyDowncastError, PyErr, PyNativeType, PyObject, PyResult, PyTryFrom, + Python, +}; +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// A Python object returned by `__await__`. +/// +/// # Examples +/// +/// ```rust +/// use pyo3::prelude::*; +/// +/// # fn main() -> PyResult<()> { +/// Python::with_gil(|py| -> PyResult<()> { +/// let list = py.eval("iter([1, 2, 3, 4])", None, None)?; +/// let numbers: PyResult> = list +/// .iter()? +/// .map(|i| i.and_then(PyAny::extract::)) +/// .collect(); +/// let sum: usize = numbers?.iter().sum(); +/// assert_eq!(sum, 10); +/// Ok(()) +/// }) +/// # } +/// ``` +#[repr(transparent)] +pub struct PyFuture(PyAny); +pyobject_native_type_named!(PyFuture); +pyobject_native_type_extract!(PyFuture); + +fn is_awaitable(obj: &Py2<'_, PyAny>) -> PyResult { + static IS_AWAITABLE: GILOnceCell = GILOnceCell::new(); + let import = || PyResult::Ok(obj.py().import("inspect")?.getattr("isawaitable")?.into()); + IS_AWAITABLE + .get_or_try_init(obj.py(), import)? + .call1(obj.py(), (obj,))? + .extract(obj.py()) +} + +impl PyFuture { + /// Constructs a `PyFuture` from a Python awaitable object. + /// + /// Equivalent to calling `__await__` method (or `__iter__` for generator-based coroutines). + pub fn from_object(obj: &PyAny) -> PyResult<&PyFuture> { + Self::from_object2(Py2::borrowed_from_gil_ref(&obj)).map(|py2| { + // Can't use into_gil_ref here because T: PyTypeInfo bound is not satisfied + // Safety: into_ptr produces a valid pointer to PyFuture object + unsafe { obj.py().from_owned_ptr(py2.into_ptr()) } + }) + } + + pub(crate) fn from_object2<'py>(obj: &Py2<'py, PyAny>) -> PyResult> { + let __await__ = intern!(obj.py(), "__await__"); + match obj.call_method0(__await__) { + Ok(obj) => Ok(unsafe { obj.downcast_into_unchecked() }), + Err(err) if err.is_instance_of::(obj.py()) => { + if obj.hasattr(__await__)? { + Err(err) + } else if is_awaitable(obj)? { + unsafe { + Py2::from_owned_ptr_or_err(obj.py(), ffi::PyObject_GetIter(obj.as_ptr())) + .map(|any| any.downcast_into_unchecked()) + } + } else { + Err(PyTypeError::new_err(format!( + "object {tp} can't be used in 'await' expression", + tp = obj.get_type().name()? + ))) + } + } + Err(err) => Err(err), + } + } + + pub(crate) fn next( + &self, + prev_result: &Option>, + ) -> PyResult> { + let py = self.0.py(); + match prev_result { + Some(Ok(val)) => { + cfg_if::cfg_if! { + if #[cfg(all(Py_3_10, not(PyPy), not(Py_LIMITED_API)))] { + let mut result = std::ptr::null_mut(); + match unsafe { ffi::PyIter_Send(self.0.as_ptr(), val.as_ptr(), &mut result) } + { + -1 => Err(PyErr::take(py).unwrap()), + 0 => Ok(IterNextOutput::Return(unsafe { + PyObject::from_owned_ptr(py, result) + })), + 1 => Ok(IterNextOutput::Yield(unsafe { + PyObject::from_owned_ptr(py, result) + })), + _ => unreachable!(), + } + } else { + let send = intern!(py, "send"); + if val.is_none(py) || !self.0.hasattr(send).unwrap_or(false) { + self.0.call_method0(intern!(py, "__next__")) + } else { + self.0.call_method1(send, (val,)) + } + .map(Into::into) + .map(IterNextOutput::Yield) + } + } + } + Some(Err(err)) => { + let throw = intern!(py, "throw"); + if self.0.hasattr(throw).unwrap_or(false) { + self.0 + .call_method1(throw, (err,)) + .map(Into::into) + .map(IterNextOutput::Yield) + } else { + Err(PyErr::from_value(err.as_ref(py))) + } + } + None => { + let close = intern!(py, "close"); + if self.0.hasattr(close).unwrap_or(false) { + self.0 + .call_method0(close) + .map(Into::into) + .map(IterNextOutput::Return) + } else { + Ok(IterNextOutput::Return(py.None().into())) + } + } + } + } +} + +impl Future for Py { + type Output = PyResult; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + FUTURE_OR_POLL.set(Some(FutureOrPoll::Future(self.clone()))); + cx.waker().wake_by_ref(); + match FUTURE_OR_POLL.take() { + Some(FutureOrPoll::Poll(poll)) => poll, + Some(FutureOrPoll::Future(_)) => { + panic!("PyFuture must be awaited in coroutine context") + } + None => unreachable!(), + } + } +} + +impl<'v> PyTryFrom<'v> for PyFuture { + fn try_from>(value: V) -> Result<&'v PyFuture, PyDowncastError<'v>> { + Err(PyDowncastError::new(value.into(), "Future")) + } + + fn try_from_exact>(value: V) -> Result<&'v PyFuture, PyDowncastError<'v>> { + value.into().downcast() + } + + #[inline] + unsafe fn try_from_unchecked>(value: V) -> &'v PyFuture { + let ptr = value.into() as *const _ as *const PyFuture; + &*ptr + } +} + +impl Py { + /// Constructs a `PyFuture` from a Python awaitable object. + /// + /// Equivalent to calling `__await__` method (or `__iter__` for generator-based coroutines). + pub fn from_object(py: Python<'_>, awaitable: PyObject) -> PyResult { + Ok(PyFuture::from_object(awaitable.as_ref(py))?.into_py(py)) + } + /// Borrows a GIL-bound reference to the PyFuture. By binding to the GIL lifetime, this + /// allows the GIL-bound reference to not require `Python` for any of its methods. + pub fn as_ref<'py>(&'py self, _py: Python<'py>) -> &'py PyFuture { + let any = self.as_ptr() as *const PyAny; + unsafe { PyNativeType::unchecked_downcast(&*any) } + } + + /// Similar to [`as_ref`](#method.as_ref), and also consumes this `Py` and registers the + /// Python object reference in PyO3's object storage. The reference count for the Python + /// object will not be decreased until the GIL lifetime ends. + pub fn into_ref(self, py: Python<'_>) -> &PyFuture { + unsafe { py.from_owned_ptr(self.into_ptr()) } + } +} diff --git a/src/types/mod.rs b/src/types/mod.rs index c6bc1972991..df8414f9dbe 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -24,6 +24,7 @@ pub use self::frozenset::{PyFrozenSet, PyFrozenSetBuilder}; pub use self::function::PyCFunction; #[cfg(all(not(Py_LIMITED_API), not(PyPy)))] pub use self::function::PyFunction; +pub use self::future::PyFuture; pub use self::iterator::PyIterator; pub use self::list::PyList; pub use self::mapping::PyMapping; @@ -287,6 +288,7 @@ mod floatob; mod frame; mod frozenset; mod function; +mod future; mod iterator; pub(crate) mod list; mod mapping; diff --git a/tests/test_awaitable.rs b/tests/test_awaitable.rs new file mode 100644 index 00000000000..229fab1ba51 --- /dev/null +++ b/tests/test_awaitable.rs @@ -0,0 +1,134 @@ +#![cfg(feature = "macros")] + +use std::task::Poll; + +use futures::{future::poll_fn, FutureExt}; +use pyo3::{prelude::*, py_run, types::PyFuture}; + +#[path = "../src/tests/common.rs"] +mod common; + +#[pyfunction] +async fn wrap_awaitable(awaitable: PyObject) -> PyResult { + let future: Py = Python::with_gil(|gil| Py::from_object(gil, awaitable))?; + future.await +} + +#[test] +fn awaitable() { + Python::with_gil(|gil| { + let wrap_awaitable = wrap_pyfunction!(wrap_awaitable, gil).unwrap(); + let test = r#" + import asyncio; + async def main(): + return await wrap_awaitable(asyncio.sleep(0.001, 42)) + assert asyncio.run(main()) == 42 + "#; + let globals = gil.import("__main__").unwrap().dict(); + globals.set_item("wrap_awaitable", wrap_awaitable).unwrap(); + gil.run(&pyo3::unindent::unindent(test), Some(globals), None) + .unwrap(); + }) +} + +#[test] +#[should_panic(expected = "object int can't be used in 'await' expression")] +fn pyfuture_not_awaitable() { + Python::with_gil(|gil| { + let wrap_awaitable = wrap_pyfunction!(wrap_awaitable, gil).unwrap(); + let test = r#" + import asyncio; + async def main(): + return await wrap_awaitable(42) + asyncio.run(main()) + "#; + let globals = gil.import("__main__").unwrap().dict(); + globals.set_item("wrap_awaitable", wrap_awaitable).unwrap(); + gil.run(&pyo3::unindent::unindent(test), Some(globals), None) + .unwrap(); + }) +} + +#[test] +#[should_panic(expected = "PyFuture must be awaited in coroutine context")] +fn pyfuture_without_coroutine() { + #[pyfunction] + fn block_on(awaitable: PyObject) -> PyResult { + let future: Py = Python::with_gil(|gil| Py::from_object(gil, awaitable))?; + futures::executor::block_on(future) + } + Python::with_gil(|gil| { + let block_on = wrap_pyfunction!(block_on, gil).unwrap(); + let test = r#" + async def coro(): + ... + block_on(coro()) + "#; + py_run!(gil, block_on, test); + }) +} + +async fn checkpoint() { + let mut ready = false; + poll_fn(|cx| { + if ready { + return Poll::Ready(()); + } + ready = true; + cx.waker().wake_by_ref(); + Poll::Pending + }) + .await +} + +#[test] +#[should_panic(expected = "Python awaitable mixed with Rust future")] +fn pyfuture_in_select() { + #[pyfunction] + async fn select(awaitable: PyObject) -> PyResult { + let future: Py = Python::with_gil(|gil| Py::from_object(gil, awaitable))?; + futures::select_biased! { + _ = checkpoint().fuse() => unreachable!(), + res = future.fuse() => res, + } + } + Python::with_gil(|gil| { + let select = wrap_pyfunction!(select, gil).unwrap(); + let test = r#" + import asyncio; + async def main(): + return await select(asyncio.sleep(1)) + asyncio.run(main()) + "#; + let globals = gil.import("__main__").unwrap().dict(); + globals.set_item("select", select).unwrap(); + gil.run(&pyo3::unindent::unindent(test), Some(globals), None) + .unwrap(); + }) +} + +#[test] +#[should_panic(expected = "Python awaitable mixed with Rust future")] +fn pyfuture_in_select2() { + #[pyfunction] + async fn select2(awaitable: PyObject) -> PyResult { + let future: Py = Python::with_gil(|gil| Py::from_object(gil, awaitable))?; + futures::select_biased! { + res = future.fuse() => res, + _ = checkpoint().fuse() => unreachable!(), + } + } + Python::with_gil(|gil| { + let select2 = wrap_pyfunction!(select2, gil).unwrap(); + let test = r#" + import asyncio; + async def main(): + return await select2(asyncio.sleep(1)) + asyncio.run(main()) + "#; + let globals = gil.import("__main__").unwrap().dict(); + globals.set_item("select2", select2).unwrap(); + gil.run(&pyo3::unindent::unindent(test), Some(globals), None) + .unwrap(); + }) +} diff --git a/tests/test_coroutine.rs b/tests/test_coroutine.rs index 8eba40787d7..41650f9a822 100644 --- a/tests/test_coroutine.rs +++ b/tests/test_coroutine.rs @@ -28,7 +28,7 @@ fn noop_coroutine() { } Python::with_gil(|gil| { let noop = wrap_pyfunction!(noop, gil).unwrap(); - let test = "import asyncio; assert asyncio.run(noop()) == 42"; + let test = "import asyncio; print(asyncio.run(noop())); assert asyncio.run(noop()) == 42"; py_run!(gil, noop, &handle_windows(test)); }) }