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
55
56
57
58
59
60
61
62
63
64
65
66

//! Macros for poem

#![forbid(unsafe_code)]
#![deny(private_in_public, unreachable_pub)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![warn(missing_docs)]

mod utils;

use proc_macro2::TokenStream;
use quote::quote;
use syn::{Error, FnArg, ItemFn, Result};

/// Wrap an asynchronous function as an `Endpoint`.
#[proc_macro_attribute]
pub fn handler(
    args: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    match generate_handler(args.into(), input.into()) {
        Ok(stream) => stream.into(),
        Err(err) => err.into_compile_error().into(),
    }
}

fn generate_handler(_args: TokenStream, input: TokenStream) -> Result<TokenStream> {
    let crate_name = utils::get_crate_name();
    let item_fn = syn::parse2::<ItemFn>(input)?;
    let vis = &item_fn.vis;
    let ident = &item_fn.sig.ident;

    if item_fn.sig.asyncness.is_none() {
        return Err(Error::new_spanned(&item_fn, "must be asynchronous"));
    }

    let mut extractors = Vec::new();
    let mut args = Vec::new();
    for input in &item_fn.sig.inputs {
        if let FnArg::Typed(pat) = input {
            let ty = &pat.ty;
            let pat = &pat.pat;
            args.push(pat);
            extractors.push(quote! { let #pat = <#ty as #crate_name::FromRequest>::from_request(&req, &mut body).await?; });
        }
    }

    let expanded = quote! {
        #[allow(non_camel_case_types)]
        #vis struct #ident;

        #[#crate_name::async_trait]
        impl #crate_name::Endpoint for #ident {
            async fn call(
                &self,
                mut req: #crate_name::Request,
            ) -> #crate_name::Result<#crate_name::Response> {
                let mut body = req.take_body().ok();
                #(#extractors)*
                #item_fn
                #crate_name::IntoResponse::into_response(#ident(#(#args),*).await)
            }
        }
    };

    Ok(expanded)
}