Supporting async for #[pyfunction] and #[pymethods]

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",

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,

pyo3-macros-backend/src/module.rs

@@ -180,6 +180,7 @@ pub fn add_fn_to_module(
     let function_wrapper_ident = function_wrapper_ident(&func.sig.ident);
 
     let spec = method::FnSpec {
+        is_async: func.sig.asyncness.is_some(),
         tp: method::FnType::FnStatic,
         name: &function_wrapper_ident,
         python_name,

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

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<PyObject> {
+    async_std::task::sleep(Duration::from_secs(1)).await;
+    Ok(Python::with_gil(|py| py.None()))
+}
+
+#[pyfunction]
+async fn sleep_for(duration: PyObject) -> PyResult<PyObject> {
+    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<PyObject> {
+    //     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"
+);