Added basic tokio main proc macro, added some docs for the attributes feature

Author: Andrew J Westlake

Cargo.toml

@@ -25,7 +25,12 @@ default = []
 [[example]]
 name = "async_std"
 path = "examples/async_std.rs"
-required-features = ["async-std-runtime"]
+required-features = ["attributes", "async-std-runtime"]
+
+[[example]]
+name = "tokio"
+path = "examples/tokio.rs"
+required-features = ["attributes", "tokio-runtime"]
 
 [[test]]
 name = "test_async_std_asyncio"

examples/tokio.rs

@@ -0,0 +1,15 @@
+#[pyo3_asyncio::tokio::main]
+async fn main() -> PyResult<()> {
+    let fut = Python::with_gil(|py| {
+        let asyncio = py.import("asyncio")?;
+
+        // convert asyncio.sleep into a Rust Future
+        pyo3_asyncio::into_future(asyncio.call_method1("sleep", (1.into_py(py),))?)
+    })?;
+
+    println!("sleeping for 1s");
+    fut.await?;
+    println!("done");
+
+    Ok(())
+}

pyo3-asyncio-macros/src/lib.rs

@@ -6,7 +6,7 @@ use proc_macro::TokenStream;
 use quote::{quote, quote_spanned};
 use syn::spanned::Spanned;
 
-/// Enables an async main function.
+/// Enables an async main function that uses the async-std runtime.
 ///
 /// # Examples
 ///
@@ -66,3 +66,65 @@ pub fn async_std_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
 
     result.into()
 }
+
+/// Enables an async main function that uses the tokio runtime.
+///
+/// # Examples
+///
+/// ```ignore
+/// #[pyo3_asyncio::tokio::main]
+/// async fn main() -> PyResult<()> {
+///     Ok(())
+/// }
+/// ```
+#[cfg(not(test))] // NOTE: exporting main breaks tests, we should file an issue.
+#[proc_macro_attribute]
+pub fn tokio_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = syn::parse_macro_input!(item as syn::ItemFn);
+
+    let ret = &input.sig.output;
+    let inputs = &input.sig.inputs;
+    let name = &input.sig.ident;
+    let body = &input.block;
+    let attrs = &input.attrs;
+    let vis = &input.vis;
+
+    if name != "main" {
+        return TokenStream::from(quote_spanned! { name.span() =>
+            compile_error!("only the main function can be tagged with #[async_std::main]"),
+        });
+    }
+
+    if input.sig.asyncness.is_none() {
+        return TokenStream::from(quote_spanned! { input.span() =>
+            compile_error!("the async keyword is missing from the function declaration"),
+        });
+    }
+
+    let result = quote! {
+        #vis fn main() {
+            #(#attrs)*
+            async fn main(#inputs) #ret {
+                #body
+            }
+
+            use pyo3::prelude::*;
+
+            pyo3_asyncio::tokio::init_multi_thread();
+
+            Python::with_gil(|py| {
+                pyo3_asyncio::with_runtime(py, || {
+                    pyo3_asyncio::tokio::run_until_complete(py, main())?;
+
+                    Ok(())
+                })
+                .map_err(|e| {
+                    e.print_and_set_sys_last_vars(py);
+                })
+                .unwrap();
+            });
+        }
+    };
+
+    result.into()
+}

src/async_std.rs

@@ -5,6 +5,7 @@ use pyo3::prelude::*;
 
 use crate::generic::{self, JoinError, Runtime};
 
+/// <span class="module-item stab portability" style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"><code>attributes</code></span> Sets up the async-std runtime and runs an async fn as main
 #[cfg(feature = "attributes")]
 pub use pyo3_asyncio_macros::async_std_main as main;
 

src/lib.rs

@@ -40,6 +40,19 @@
 //! <span
 //!   class="module-item stab portability"
 //!   style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"
+//! ><code>attributes</code></span>
+//! are only available when the `attributes` Cargo feature is enabled:
+//!
+//! ```toml
+//! [dependencies.pyo3-asyncio]
+//! version = "0.13.0"
+//! features = ["attributes"]
+//! ```
+//!
+//! Items marked with
+//! <span
+//!   class="module-item stab portability"
+//!   style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"
 //! ><code>async-std-runtime</code></span>
 //! are only available when the `async-std-runtime` Cargo feature is enabled:
 //!

src/tokio.rs

@@ -12,6 +12,10 @@ use pyo3::prelude::*;
 
 use crate::generic;
 
+/// <span class="module-item stab portability" style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"><code>attributes</code></span> Sets up the tokio runtime and runs an async fn as main
+#[cfg(feature = "attributes")]
+pub use pyo3_asyncio_macros::tokio_main as main;
+
 static TOKIO_RUNTIME: OnceCell<Runtime> = OnceCell::new();
 
 const EXPECT_TOKIO_INIT: &str = "Tokio runtime must be initialized";