From 2cbd2e9332e09a9b13ffe04d479288e32ec9f4b5 Mon Sep 17 00:00:00 2001 From: Andrew J Westlake Date: Wed, 27 Jan 2021 16:36:51 -0600 Subject: [PATCH 1/2] Added support for async in #[pyfunction] --- Cargo.toml | 10 +++ pyo3-macros-backend/src/method.rs | 2 + pyo3-macros-backend/src/module.rs | 11 ++- pyo3-macros-backend/src/pymethod.rs | 20 +++++ pytests/test_async_fn.rs | 117 ++++++++++++++++++++++++++++ 5 files changed, 158 insertions(+), 2 deletions(-) create mode 100644 pytests/test_async_fn.rs diff --git a/Cargo.toml b/Cargo.toml index b612b2ca65d..caeb9c06d9d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,13 +31,18 @@ serde = {version = "1.0", optional = true} [dev-dependencies] assert_approx_eq = "1.1.0" +async-std = "1.9" trybuild = "1.0.23" rustversion = "1.0" proptest = { version = "0.10.1", default-features = false, features = ["std"] } # features needed to run the PyO3 test suite pyo3 = { path = ".", default-features = false, features = ["macros", "auto-initialize"] } +pyo3-asyncio = { git = "https://github.com/awestlake87/pyo3-asyncio", branch = "attributes", features = ["attributes", "testing", "async-std-runtime"] } serde_json = "1.0.61" +[patch.crates-io] +pyo3 = { path = ".", default-features = false, features = ["macros", "auto-initialize"] } + [features] default = ["macros", "auto-initialize"] @@ -65,6 +70,11 @@ auto-initialize = [] # Optimizes PyObject to Vec conversion and so on. nightly = [] +[[test]] +name = "test_async_fn" +path = "pytests/test_async_fn.rs" +harness = false + [workspace] members = [ "pyo3-macros", diff --git a/pyo3-macros-backend/src/method.rs b/pyo3-macros-backend/src/method.rs index 2575385b2f2..a0fc5a1c53b 100644 --- a/pyo3-macros-backend/src/method.rs +++ b/pyo3-macros-backend/src/method.rs @@ -85,6 +85,7 @@ impl SelfType { #[derive(Clone, Debug)] pub struct FnSpec<'a> { + pub is_async: bool, pub tp: FnType, // Rust function name pub name: &'a syn::Ident, @@ -244,6 +245,7 @@ impl<'a> FnSpec<'a> { let doc = utils::get_doc(&meth_attrs, text_signature, true)?; Ok(FnSpec { + is_async: sig.asyncness.is_some(), tp: fn_type, name, python_name, diff --git a/pyo3-macros-backend/src/module.rs b/pyo3-macros-backend/src/module.rs index 087ce2307b5..67d1970ea27 100644 --- a/pyo3-macros-backend/src/module.rs +++ b/pyo3-macros-backend/src/module.rs @@ -58,7 +58,11 @@ pub fn process_functions_in_module(func: &mut syn::ItemFn) -> syn::Result<()> { } /// Transforms a rust fn arg parsed with syn into a method::FnArg -fn wrap_fn_argument(cap: &syn::PatType) -> syn::Result { +fn wrap_fn_argument(cap: &syn::PatType, is_async: bool) -> syn::Result { + if is_async && utils::is_python(&cap.ty) { + bail_spanned!(cap.ty.span() => "Python argument is not supported on async functions"); + } + let (mutability, by_ref, ident) = match &*cap.pat { syn::Pat::Ident(patid) => (&patid.mutability, &patid.by_ref, &patid.ident), _ => bail_spanned!(cap.pat.span() => "unsupported argument"), @@ -140,6 +144,8 @@ pub fn add_fn_to_module( python_name: Ident, pyfn_attrs: PyFunctionAttr, ) -> syn::Result { + let is_async = func.sig.asyncness.is_some(); + let mut arguments = Vec::new(); for (i, input) in func.sig.inputs.iter().enumerate() { @@ -166,7 +172,7 @@ pub fn add_fn_to_module( cap.span() => "expected &PyModule as first argument with `pass_module`" ); } else { - arguments.push(wrap_fn_argument(cap)?); + arguments.push(wrap_fn_argument(cap, is_async)?); } } } @@ -180,6 +186,7 @@ pub fn add_fn_to_module( let function_wrapper_ident = function_wrapper_ident(&func.sig.ident); let spec = method::FnSpec { + is_async, tp: method::FnType::FnStatic, name: &function_wrapper_ident, python_name, diff --git a/pyo3-macros-backend/src/pymethod.rs b/pyo3-macros-backend/src/pymethod.rs index d6234cdadde..599861acca2 100644 --- a/pyo3-macros-backend/src/pymethod.rs +++ b/pyo3-macros-backend/src/pymethod.rs @@ -91,6 +91,16 @@ fn impl_wrap_common( slf: TokenStream, body: TokenStream, ) -> TokenStream { + let body = if spec.is_async { + quote! { + pyo3::Python::with_gil(move |py| pyo3_asyncio::async_std::into_coroutine(py, #body)) + } + } else { + quote! { + #body + } + }; + let python_name = &spec.python_name; if spec.args.is_empty() && noargs { quote! { @@ -380,6 +390,16 @@ pub fn impl_arg_params( self_: Option<&syn::Type>, body: TokenStream, ) -> TokenStream { + let body = if spec.is_async { + quote! { + pyo3::Python::with_gil(move |py| pyo3_asyncio::async_std::into_coroutine(py, #body)) + } + } else { + quote! { + #body + } + }; + if spec.args.is_empty() { return quote! { #body diff --git a/pytests/test_async_fn.rs b/pytests/test_async_fn.rs new file mode 100644 index 00000000000..172973a1e08 --- /dev/null +++ b/pytests/test_async_fn.rs @@ -0,0 +1,117 @@ +use std::time::Duration; + +use pyo3::prelude::*; + +const TEST_MOD: &'static str = r#" +import asyncio + +async def py_sleep(duration): + await asyncio.sleep(duration) + +async def sleep(sleeper): + await sleeper() + +async def sleep_for(sleeper, duration): + await sleeper(duration) +"#; + +#[pyfunction] +async fn sleep() -> PyResult { + async_std::task::sleep(Duration::from_secs(1)).await; + Ok(Python::with_gil(|py| py.None())) +} + +#[pyfunction] +async fn sleep_for(duration: PyObject) -> PyResult { + let duration: f64 = Python::with_gil(|py| duration.as_ref(py).extract())?; + let microseconds = duration * 1.0e6; + + async_std::task::sleep(Duration::from_micros(microseconds as u64)).await; + + Ok(Python::with_gil(|py| py.None())) +} + +#[pyclass] +struct Sleeper { + duration: Duration, +} + +#[pymethods] +impl Sleeper { + // FIXME: &self screws up the 'static requirement for into_coroutine. Could be fixed by + // supporting impl Future along with async (which would be nice anyway). I don't think any + // async method that accesses member variables can be reasonably supported with the async fn + // syntax because of the 'static lifetime requirement, so it would have to fall back to + // impl Future in nearly all cases + // + // async fn sleep(&self) -> PyResult { + // let duration = self.duration.clone(); + + // async_std::task::sleep(duration).await; + + // Ok(Python::with_gil(|py| py.None())) + // } +} + +#[pyo3_asyncio::async_std::test] +async fn test_sleep() -> PyResult<()> { + let fut = Python::with_gil(|py| { + let sleeper_mod = PyModule::new(py, "rust_sleeper")?; + sleeper_mod.add_wrapped(pyo3::wrap_pyfunction!(sleep))?; + + let test_mod = + PyModule::from_code(py, TEST_MOD, "test_rust_coroutine/test_mod.py", "test_mod")?; + + pyo3_asyncio::into_future(test_mod.call_method1("sleep", (sleeper_mod.getattr("sleep")?,))?) + })?; + + fut.await?; + + Ok(()) +} + +#[pyo3_asyncio::async_std::test] +async fn test_sleep_for() -> PyResult<()> { + let fut = Python::with_gil(|py| { + let sleeper_mod = PyModule::new(py, "rust_sleeper")?; + sleeper_mod.add_wrapped(pyo3::wrap_pyfunction!(sleep_for))?; + + let test_mod = + PyModule::from_code(py, TEST_MOD, "test_rust_coroutine/test_mod.py", "test_mod")?; + + pyo3_asyncio::into_future(test_mod.call_method1( + "sleep_for", + (sleeper_mod.getattr("sleep_for")?, 2.into_py(py)), + )?) + })?; + + fut.await?; + + Ok(()) +} + +// #[pyo3_asyncio::async_std::test] +// async fn test_sleeper() -> PyResult<()> { +// let fut = Python::with_gil(|py| { +// let sleeper = PyCell::new( +// py, +// Sleeper { +// duration: Duration::from_secs(3), +// }, +// )?; + +// let test_mod = +// PyModule::from_code(py, TEST_MOD, "test_rust_coroutine/test_mod.py", "test_mod")?; + +// pyo3_asyncio::into_future(test_mod.call_method1("sleep_for", (sleeper.getattr("sleep")?,))?) +// })?; + +// fut.await?; + +// Ok(()) +// } + +pyo3_asyncio::testing::test_main!( + #[pyo3_asyncio::async_std::main], + "Async #[pyfunction] Test Suite" +); From d604b77a2fbc63b421c6e8099b6ed5aeffead0f1 Mon Sep 17 00:00:00 2001 From: Andrew J Westlake Date: Wed, 27 Jan 2021 16:46:44 -0600 Subject: [PATCH 2/2] Reversed restriction on Python argument since it can be addressed by 'static and Send requirements for into_coroutine --- pyo3-macros-backend/src/module.rs | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/pyo3-macros-backend/src/module.rs b/pyo3-macros-backend/src/module.rs index 67d1970ea27..bd850a64eb5 100644 --- a/pyo3-macros-backend/src/module.rs +++ b/pyo3-macros-backend/src/module.rs @@ -58,11 +58,7 @@ pub fn process_functions_in_module(func: &mut syn::ItemFn) -> syn::Result<()> { } /// Transforms a rust fn arg parsed with syn into a method::FnArg -fn wrap_fn_argument(cap: &syn::PatType, is_async: bool) -> syn::Result { - if is_async && utils::is_python(&cap.ty) { - bail_spanned!(cap.ty.span() => "Python argument is not supported on async functions"); - } - +fn wrap_fn_argument(cap: &syn::PatType) -> syn::Result { let (mutability, by_ref, ident) = match &*cap.pat { syn::Pat::Ident(patid) => (&patid.mutability, &patid.by_ref, &patid.ident), _ => bail_spanned!(cap.pat.span() => "unsupported argument"), @@ -144,8 +140,6 @@ pub fn add_fn_to_module( python_name: Ident, pyfn_attrs: PyFunctionAttr, ) -> syn::Result { - let is_async = func.sig.asyncness.is_some(); - let mut arguments = Vec::new(); for (i, input) in func.sig.inputs.iter().enumerate() { @@ -172,7 +166,7 @@ pub fn add_fn_to_module( cap.span() => "expected &PyModule as first argument with `pass_module`" ); } else { - arguments.push(wrap_fn_argument(cap, is_async)?); + arguments.push(wrap_fn_argument(cap)?); } } } @@ -186,7 +180,7 @@ pub fn add_fn_to_module( let function_wrapper_ident = function_wrapper_ident(&func.sig.ident); let spec = method::FnSpec { - is_async, + is_async: func.sig.asyncness.is_some(), tp: method::FnType::FnStatic, name: &function_wrapper_ident, python_name,