mkutils-macros 0.1.196

Utility methods, traits, and types.
Documentation
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 MethodApplications {
    method_applications: Vec<TokenStream2>,
}

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

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

        Ok(method_applications)
    }
}

pub fn tokio_main(attr_args_token_stream: TokenStream, input_token_stream: TokenStream) -> TokenStream {
    let MethodApplications { method_applications } = 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()
                #(#method_applications)*
                .build()
                .unwrap()
                .block_on(async #block)
        }
    };

    tokio_main_fn_token_stream.into()
}