Added config knobs to pyo3_asyncio::tokio::main and examples for each

Author: Andrew J Westlake

Cargo.toml

@@ -32,6 +32,16 @@ name = "tokio"
 path = "examples/tokio.rs"
 required-features = ["attributes", "tokio-runtime"]
 
+[[example]]
+name = "tokio_current_thread"
+path = "examples/tokio_current_thread.rs"
+required-features = ["attributes", "tokio-runtime"]
+
+[[example]]
+name = "tokio_multi_thread"
+path = "examples/tokio_multi_thread.rs"
+required-features = ["attributes", "tokio-runtime"]
+
 [[test]]
 name = "test_async_std_asyncio"
 path = "pytests/test_async_std_asyncio.rs"

examples/async_std.rs

@@ -1,3 +1,5 @@
+use pyo3::prelude::*;
+
 #[pyo3_asyncio::async_std::main]
 async fn main() -> PyResult<()> {
     let fut = Python::with_gil(|py| {

examples/tokio.rs

@@ -1,3 +1,5 @@
+use pyo3::prelude::*;
+
 #[pyo3_asyncio::tokio::main]
 async fn main() -> PyResult<()> {
     let fut = Python::with_gil(|py| {

examples/tokio_current_thread.rs

@@ -0,0 +1,17 @@
+use pyo3::prelude::*;
+
+#[pyo3_asyncio::tokio::main(flavor = "current_thread")]
+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(())
+}

examples/tokio_multi_thread.rs

@@ -0,0 +1,17 @@
+use pyo3::prelude::*;
+
+#[pyo3_asyncio::tokio::main(flavor = "multi_thread", worker_threads = 10)]
+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/Cargo.toml

@@ -8,5 +8,6 @@ edition = "2018"
 proc-macro = true
 
 [dependencies]
+proc-macro2 = "1.0"
 quote = "1"
 syn = { version = "1", features = ["full"] }

pyo3-asyncio-macros/src/lib.rs

@@ -2,6 +2,8 @@
 #![deny(missing_debug_implementations, nonstandard_style)]
 #![recursion_limit = "512"]
 
+mod tokio;
+
 use proc_macro::TokenStream;
 use quote::{quote, quote_spanned};
 use syn::spanned::Spanned;
@@ -47,9 +49,7 @@ pub fn async_std_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
                 #body
             }
 
-            use pyo3::prelude::*;
-
-            Python::with_gil(|py| {
+            pyo3::Python::with_gil(|py| {
                 pyo3_asyncio::with_runtime(py, || {
                     pyo3_asyncio::async_std::run_until_complete(py, main())?;
 
@@ -61,70 +61,44 @@ pub fn async_std_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
                 .unwrap();
             });
         }
-
     };
 
     result.into()
 }
 
 /// Enables an async main function that uses the tokio runtime.
 ///
+/// # Arguments
+/// * `flavor` - selects the type of tokio runtime ["multi_thread", "current_thread"]
+/// * `worker_threads` - number of worker threads, defaults to the number of CPUs on the system
+///
 /// # Examples
 ///
+/// Default configuration:
 /// ```ignore
 /// #[pyo3_asyncio::tokio::main]
 /// async fn main() -> PyResult<()> {
 ///     Ok(())
 /// }
 /// ```
+///
+/// Current-thread scheduler:
+/// ```ignore
+/// #[pyo3_asyncio::tokio::main(flavor = "current_thread")]
+/// async fn main() -> PyResult<()> {
+///     Ok(())
+/// }
+/// ```
+///
+/// Multi-thread scheduler with custom worker thread count:
+/// ```ignore
+/// #[pyo3_asyncio::tokio::main(flavor = "multi_thread", worker_threads = 10)]
+/// 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()
+pub fn tokio_main(args: TokenStream, item: TokenStream) -> TokenStream {
+    tokio::main(args, item, true)
 }