tokio_extra_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, parse_quote_spanned, spanned::Spanned, ItemFn, ReturnType};
4
5#[proc_macro_attribute]
6pub fn background_task(_attr: TokenStream, tokens: TokenStream) -> TokenStream {
7    let mut item = parse_macro_input!(tokens as ItemFn);
8
9    // Change `async fn()` to `fn()`
10    item.sig.asyncness = None;
11    let ItemFn {
12        attrs,
13        vis,
14        mut sig,
15        block,
16    } = item;
17
18    // Change fn
19    match &sig.output {
20        ReturnType::Default => {}
21        ReturnType::Type(_, ret) => {
22            let ret = quote!(#ret);
23            sig.output = parse_quote_spanned! {ret.span()=>
24                -> tokio::task::JoinHandle<#ret>
25            };
26        }
27    };
28
29    let stmts = &block.stmts;
30
31    let new_tokens = match &sig.output {
32        ReturnType::Default => {
33            quote! {
34                #(#attrs)* #vis #sig {
35                    tokio::spawn(async {
36                        #(#stmts)*
37                    });
38                }
39            }
40        }
41        ReturnType::Type(_, _) => {
42            quote! {
43                #(#attrs)* #vis #sig {
44                    tokio::spawn(async {
45                        #(#stmts)*
46                    })
47                }
48            }
49        }
50    };
51
52    new_tokens.into()
53}
tokio_extra_macros/lib.rs

tokio_extra_macros/
lib.rs
