macros: Allow arguments in non-main functions

This commit is contained in:
Douman 2019-09-24 16:03:26 +02:00
parent 5efe31f2ed
commit a1d1eb5eb3
3 changed files with 24 additions and 30 deletions

View File

@ -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() {}

View File

@ -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]
| ^^^^^^^

View File

@ -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
@ -57,6 +61,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;
@ -65,7 +70,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()
@ -95,16 +100,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 })
}
},
RuntimeType::Auto => quote! {