macros: Allow arguments in non-main functions

build-tests/tests/fail/macros_invalid_input.rs

@@ -3,9 +3,6 @@ use build_tests::tokio;
 #[tokio::main]
 fn main_is_not_async() {}
 
-#[tokio::main]
-async fn main_fn_has_args(_x: u8) {}
-
 #[tokio::main(foo)]
 async fn main_attr_has_unknown_args() {}
 

build-tests/tests/fail/macros_invalid_input.stderr

@@ -4,44 +4,38 @@ error: the async keyword is missing from the function declaration
 4 | fn main_is_not_async() {}
   | ^^
 
-error: the main function cannot accept arguments
- --> $DIR/macros_invalid_input.rs:7:27
-  |
-7 | async fn main_fn_has_args(_x: u8) {}
-  |                           ^^^^^^
-
 error: Unknown attribute foo is specified
- --> $DIR/macros_invalid_input.rs:9:15
+ --> $DIR/macros_invalid_input.rs:6:15
   |
-9 | #[tokio::main(foo)]
+6 | #[tokio::main(foo)]
   |               ^^^
 
 error: Must have specified ident
-  --> $DIR/macros_invalid_input.rs:12:15
-   |
-12 | #[tokio::main(multi_thread::bar)]
-   |               ^^^^^^^^^^^^^^^^^
+ --> $DIR/macros_invalid_input.rs:9:15
+  |
+9 | #[tokio::main(multi_thread::bar)]
+  |               ^^^^^^^^^^^^^^^^^
 
 error: the async keyword is missing from the function declaration
-  --> $DIR/macros_invalid_input.rs:16:1
+  --> $DIR/macros_invalid_input.rs:13:1
    |
-16 | fn test_is_not_async() {}
+13 | fn test_is_not_async() {}
    | ^^
 
 error: the test function cannot accept arguments
-  --> $DIR/macros_invalid_input.rs:19:27
+  --> $DIR/macros_invalid_input.rs:16:27
    |
-19 | async fn test_fn_has_args(_x: u8) {}
+16 | async fn test_fn_has_args(_x: u8) {}
    |                           ^^^^^^
 
 error: unexpected token
-  --> $DIR/macros_invalid_input.rs:21:15
+  --> $DIR/macros_invalid_input.rs:18:15
    |
-21 | #[tokio::test(foo)]
+18 | #[tokio::test(foo)]
    |               ^^^
 
 error: second test attribute is supplied
-  --> $DIR/macros_invalid_input.rs:25:1
+  --> $DIR/macros_invalid_input.rs:22:1
    |
-25 | #[test]
+22 | #[test]
    | ^^^^^^^

tokio-macros/src/lib.rs

@@ -25,6 +25,10 @@ use quote::quote;
 /// - `single_thread` - Uses `current_thread`.
 /// - `multi_thread` - Uses multi-threaded runtime. Used by default.
 ///
+/// ## Function arguments:
+///
+/// Arguments are allowed for any functions aside from `main` which is special
+///
 /// ## Usage
 ///
 /// ### Select runtime
@@ -56,6 +60,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 
     let ret = &input.sig.output;
     let name = &input.sig.ident;
+    let inputs = &input.sig.inputs;
     let body = &input.block;
     let attrs = &input.attrs;
 
@@ -64,7 +69,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
         return syn::Error::new_spanned(input.sig.fn_token, msg)
             .to_compile_error()
             .into();
-    } else if !input.sig.inputs.is_empty() {
+    } else if name == "main" && !inputs.is_empty() {
         let msg = "the main function cannot accept arguments";
         return syn::Error::new_spanned(&input.sig.inputs, msg)
             .to_compile_error()
@@ -94,16 +99,14 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
     let result = match runtime {
         RuntimeType::Multi => quote! {
             #(#attrs)*
-            fn #name() #ret {
-                let mut rt = tokio::runtime::Runtime::new().unwrap();
-                rt.block_on(async { #body })
+            fn #name(#inputs) #ret {
+                tokio::runtime::Runtime::new().unwrap().block_on(async { #body })
             }
         },
         RuntimeType::Single => quote! {
             #(#attrs)*
-            fn #name() #ret {
-                let mut rt = tokio::runtime::current_thread::Runtime::new().unwrap();
-                rt.block_on(async { #body })
+            fn #name(#inputs) #ret {
+                tokio::runtime::current_thread::Runtime::new().unwrap().block_on(async { #body })
             }
         },
     };