mkutils-macros 0.1.136

Utility methods, traits, and types.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

use crate::utils::{Cat3, CommaPunctuated, IdentAssignment};
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use syn::{
    Error as SynError, Expr, ItemFn,
    parse::{Parse, ParseStream},
};

pub struct BuilderCalls {
    builder_calls: Vec<TokenStream2>,
}

impl BuilderCalls {
    fn builder_call(Cat3(ident, _comma, expr): IdentAssignment<Expr>) -> TokenStream2 {
        quote::quote! { .#ident(#expr) }
    }
}

impl Parse for BuilderCalls {
    fn parse(parse_stream: ParseStream) -> Result<Self, SynError> {
        let assignments = CommaPunctuated::<IdentAssignment<Expr>>::parse_terminated(parse_stream)?;
        let builder_calls = assignments.into_iter().map(Self::builder_call).collect();
        let builder_calls = Self { builder_calls };

        Ok(builder_calls)
    }
}

pub fn tokio_main(attr_args_token_stream: TokenStream, input_token_stream: TokenStream) -> TokenStream {
    let BuilderCalls { builder_calls } = syn::parse_macro_input!(attr_args_token_stream);
    let ItemFn {
        attrs,
        vis,
        mut sig,
        block,
    } = syn::parse_macro_input!(input_token_stream);

    sig.asyncness = None;

    // NOTE: [https://docs.rs/tokio/latest/tokio/attr.main.html#using-the-multi-threaded-runtime]
    let tokio_main_fn_token_stream = quote::quote! {
        #(#attrs)*
        #vis #sig {
            ::tokio::runtime::Builder::new_multi_thread()
                .enable_all()
                #(#builder_calls)*
                .build()
                .unwrap()
                .block_on(async #block)
        }
    };

    tokio_main_fn_token_stream.into()
}