poem_derive/
lib.rs

1//! Macros for poem
2
3#![doc(html_favicon_url = "https://raw.githubusercontent.com/poem-web/poem/master/favicon.ico")]
4#![doc(html_logo_url = "https://raw.githubusercontent.com/poem-web/poem/master/logo.png")]
5#![forbid(unsafe_code)]
6#![deny(unreachable_pub)]
7#![cfg_attr(docsrs, feature(doc_cfg))]
8#![warn(missing_docs)]
9
10mod utils;
11
12use proc_macro::TokenStream;
13use quote::{format_ident, quote};
14use syn::{FnArg, GenericParam, ItemFn, Member, Result, parse_macro_input};
15
16/// Wrap an asynchronous function as an `Endpoint`.
17///
18/// # Example
19///
20/// ```ignore
21/// #[handler]
22/// async fn example() {
23/// }
24/// ```
25#[proc_macro_attribute]
26pub fn handler(args: TokenStream, input: TokenStream) -> TokenStream {
27    let mut internal = false;
28
29    let arg_parser = syn::meta::parser(|meta| {
30        if meta.path.is_ident("internal") {
31            internal = true;
32        }
33        Ok(())
34    });
35    parse_macro_input!(args with arg_parser);
36
37    match generate_handler(internal, input) {
38        Ok(stream) => stream,
39        Err(err) => err.into_compile_error().into(),
40    }
41}
42
43fn generate_handler(internal: bool, input: TokenStream) -> Result<TokenStream> {
44    let crate_name = utils::get_crate_name(internal);
45    let item_fn = syn::parse::<ItemFn>(input)?;
46    let (impl_generics, type_generics, where_clause) = item_fn.sig.generics.split_for_impl();
47    let vis = &item_fn.vis;
48    let docs = item_fn
49        .attrs
50        .iter()
51        .filter(|attr| attr.path().is_ident("doc"))
52        .cloned()
53        .collect::<Vec<_>>();
54    let ident = &item_fn.sig.ident;
55    let call_await = if item_fn.sig.asyncness.is_some() {
56        Some(quote::quote!(.await))
57    } else {
58        None
59    };
60
61    let def_struct = if !item_fn.sig.generics.params.is_empty() {
62        let iter = item_fn
63            .sig
64            .generics
65            .params
66            .iter()
67            .filter_map(|param| match param {
68                GenericParam::Type(ty) => Some(ty),
69                _ => None,
70            })
71            .enumerate()
72            .map(|(idx, ty)| {
73                let ident = format_ident!("_mark{}", idx);
74                let ty_ident = &ty.ident;
75                (ident, ty_ident)
76            });
77
78        let struct_members = iter.clone().map(|(ident, ty_ident)| {
79            quote! { #ident: ::std::marker::PhantomData<#ty_ident> }
80        });
81
82        let default_members = iter.clone().map(|(ident, _ty_ident)| {
83            quote! { #ident: ::std::marker::PhantomData }
84        });
85
86        quote! {
87            #vis struct #ident #type_generics { #(#struct_members),*}
88            impl #type_generics ::std::default::Default for #ident #type_generics {
89                fn default() -> Self {
90                    Self { #(#default_members),* }
91                }
92            }
93        }
94    } else {
95        quote! { #vis struct #ident; }
96    };
97
98    let mut extractors = Vec::new();
99    let mut args = Vec::new();
100    for (idx, input) in item_fn.sig.inputs.clone().into_iter().enumerate() {
101        if let FnArg::Typed(pat) = input {
102            let ty = &pat.ty;
103            let id = quote::format_ident!("p{}", idx);
104            args.push(id.clone());
105            extractors.push(quote! {
106                let #id = <#ty as #crate_name::FromRequest>::from_request(&req, &mut body).await?;
107            });
108        }
109    }
110
111    let expanded = quote! {
112        #(#docs)*
113        #[allow(non_camel_case_types)]
114        #def_struct
115
116        impl #impl_generics #crate_name::Endpoint for #ident #type_generics #where_clause {
117            type Output = #crate_name::Response;
118
119            #[allow(unused_mut)]
120            async fn call(&self, mut req: #crate_name::Request) -> #crate_name::Result<Self::Output> {
121                let (req, mut body) = req.split();
122                #(#extractors)*
123                #item_fn
124                let res = #ident(#(#args),*)#call_await;
125                let res = #crate_name::error::IntoResult::into_result(res);
126                std::result::Result::map(res, #crate_name::IntoResponse::into_response)
127            }
128        }
129    };
130
131    Ok(expanded.into())
132}
133
134#[doc(hidden)]
135#[proc_macro]
136pub fn generate_implement_middlewares(_: TokenStream) -> TokenStream {
137    let mut impls = Vec::new();
138
139    for i in 2..=16 {
140        let idents = (0..i)
141            .map(|i| format_ident!("T{}", i + 1))
142            .collect::<Vec<_>>();
143        let output_type = idents.last().unwrap();
144        let first_ident = idents.first().unwrap();
145        let mut where_clauses = vec![quote! { #first_ident: Middleware<E> }];
146        let mut transforms = Vec::new();
147
148        for k in 1..i {
149            let prev_ident = &idents[k - 1];
150            let current_ident = &idents[k];
151            where_clauses.push(quote! { #current_ident: Middleware<#prev_ident::Output> });
152        }
153
154        for k in 0..i {
155            let n = Member::from(k);
156            transforms.push(quote! { let ep = self.#n.transform(ep); });
157        }
158
159        let expanded = quote! {
160            impl<E, #(#idents),*> Middleware<E> for (#(#idents),*)
161                where
162                    E: Endpoint,
163                    #(#where_clauses,)*
164            {
165                type Output = #output_type::Output;
166
167                fn transform(&self, ep: E) -> Self::Output {
168                    #(#transforms)*
169                    ep
170                }
171            }
172        };
173
174        impls.push(expanded);
175    }
176
177    quote!(#(#impls)*).into()
178}