From 568394a28e62904f8db4872da0d10e90d831ff42 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Thu, 18 Aug 2022 11:41:14 +0200 Subject: [PATCH] Support changing state type in `#[debug_handler]` (#1271) * support setting body type for #[debug_handler] * Use lookahead1 to give better errors and detect duplicate arguments * fix docs link --- axum-macros/Cargo.toml | 6 +- axum-macros/src/debug_handler.rs | 115 ++++++++++++++++-- axum-macros/src/lib.rs | 48 ++++++++ .../debug_handler/fail/duplicate_args.rs | 9 ++ .../debug_handler/fail/duplicate_args.stderr | 11 ++ .../debug_handler/fail/invalid_attrs.stderr | 2 +- .../tests/debug_handler/pass/infer_state.rs | 31 +++++ .../tests/debug_handler/pass/set_state.rs | 27 ++++ .../debug_handler/pass/state_and_body.rs | 10 ++ 9 files changed, 249 insertions(+), 10 deletions(-) create mode 100644 axum-macros/tests/debug_handler/fail/duplicate_args.rs create mode 100644 axum-macros/tests/debug_handler/fail/duplicate_args.stderr create mode 100644 axum-macros/tests/debug_handler/pass/infer_state.rs create mode 100644 axum-macros/tests/debug_handler/pass/set_state.rs create mode 100644 axum-macros/tests/debug_handler/pass/state_and_body.rs diff --git a/axum-macros/Cargo.toml b/axum-macros/Cargo.toml index 50d7aafc..89d84c64 100644 --- a/axum-macros/Cargo.toml +++ b/axum-macros/Cargo.toml @@ -18,7 +18,11 @@ proc-macro = true heck = "0.4" proc-macro2 = "1.0" quote = "1.0" -syn = { version = "1.0", features = ["full"] } +syn = { version = "1.0", features = [ + "full", + # needed for `Hash` impls + "extra-traits", +] } [dev-dependencies] axum = { path = "../axum", version = "0.5", features = ["headers"] } diff --git a/axum-macros/src/debug_handler.rs b/axum-macros/src/debug_handler.rs index 4972fc49..0fe8ad6a 100644 --- a/axum-macros/src/debug_handler.rs +++ b/axum-macros/src/debug_handler.rs @@ -1,8 +1,10 @@ +use std::collections::HashSet; + use proc_macro2::TokenStream; use quote::{format_ident, quote, quote_spanned}; use syn::{parse::Parse, spanned::Spanned, FnArg, ItemFn, Token, Type}; -pub(crate) fn expand(attr: Attrs, item_fn: ItemFn) -> TokenStream { +pub(crate) fn expand(mut attr: Attrs, item_fn: ItemFn) -> TokenStream { let check_extractor_count = check_extractor_count(&item_fn); let check_request_last_extractor = check_request_last_extractor(&item_fn); let check_path_extractor = check_path_extractor(&item_fn); @@ -12,8 +14,14 @@ pub(crate) fn expand(attr: Attrs, item_fn: ItemFn) -> TokenStream { // If the function is generic, we can't reliably check its inputs or whether the future it // returns is `Send`. Skip those checks to avoid unhelpful additional compiler errors. let check_inputs_and_future_send = if item_fn.sig.generics.params.is_empty() { + if attr.state_ty.is_none() { + attr.state_ty = state_type_from_args(&item_fn); + } + + let state_ty = attr.state_ty.unwrap_or_else(|| syn::parse_quote!(())); + let check_inputs_impls_from_request = - check_inputs_impls_from_request(&item_fn, &attr.body_ty); + check_inputs_impls_from_request(&item_fn, &attr.body_ty, state_ty); let check_future_send = check_future_send(&item_fn); quote! { @@ -39,21 +47,46 @@ pub(crate) fn expand(attr: Attrs, item_fn: ItemFn) -> TokenStream { } } +mod kw { + syn::custom_keyword!(body); + syn::custom_keyword!(state); +} + pub(crate) struct Attrs { body_ty: Type, + state_ty: Option, } impl Parse for Attrs { fn parse(input: syn::parse::ParseStream) -> syn::Result { let mut body_ty = None; + let mut state_ty = None; while !input.is_empty() { - let ident = input.parse::()?; - if ident == "body" { + let lh = input.lookahead1(); + + if lh.peek(kw::body) { + let kw = input.parse::()?; + if body_ty.is_some() { + return Err(syn::Error::new_spanned( + kw, + "`body` specified more than once", + )); + } input.parse::()?; body_ty = Some(input.parse()?); + } else if lh.peek(kw::state) { + let kw = input.parse::()?; + if state_ty.is_some() { + return Err(syn::Error::new_spanned( + kw, + "`state` specified more than once", + )); + } + input.parse::()?; + state_ty = Some(input.parse()?); } else { - return Err(syn::Error::new_spanned(ident, "unknown argument")); + return Err(lh.error()); } let _ = input.parse::(); @@ -61,7 +94,7 @@ impl Parse for Attrs { let body_ty = body_ty.unwrap_or_else(|| syn::parse_quote!(axum::body::Body)); - Ok(Self { body_ty }) + Ok(Self { body_ty, state_ty }) } } @@ -167,7 +200,11 @@ fn check_multiple_body_extractors(item_fn: &ItemFn) -> TokenStream { } } -fn check_inputs_impls_from_request(item_fn: &ItemFn, body_ty: &Type) -> TokenStream { +fn check_inputs_impls_from_request( + item_fn: &ItemFn, + body_ty: &Type, + state_ty: Type, +) -> TokenStream { item_fn .sig .inputs @@ -203,7 +240,7 @@ fn check_inputs_impls_from_request(item_fn: &ItemFn, body_ty: &Type) -> TokenStr #[allow(warnings)] fn #name() where - #ty: ::axum::extract::FromRequest<(), #body_ty> + Send, + #ty: ::axum::extract::FromRequest<#state_ty, #body_ty> + Send, {} } }) @@ -371,6 +408,68 @@ fn self_receiver(item_fn: &ItemFn) -> Option { None } +/// Given a signature like +/// +/// ```skip +/// #[debug_handler] +/// async fn handler( +/// _: axum::extract::State, +/// _: State, +/// ) {} +/// ``` +/// +/// This will extract `AppState`. +/// +/// Returns `None` if there are no `State` args or multiple of different types. +fn state_type_from_args(item_fn: &ItemFn) -> Option { + let state_inputs = item_fn + .sig + .inputs + .iter() + .filter_map(|input| match input { + FnArg::Receiver(_) => None, + FnArg::Typed(pat_type) => Some(pat_type), + }) + .map(|pat_type| &pat_type.ty) + .filter_map(|ty| { + if let Type::Path(path) = &**ty { + Some(&path.path) + } else { + None + } + }) + .filter_map(|path| { + if let Some(last_segment) = path.segments.last() { + if last_segment.ident != "State" { + return None; + } + + match &last_segment.arguments { + syn::PathArguments::AngleBracketed(args) if args.args.len() == 1 => { + Some(args.args.first().unwrap()) + } + _ => None, + } + } else { + None + } + }) + .filter_map(|generic_arg| { + if let syn::GenericArgument::Type(ty) = generic_arg { + Some(ty) + } else { + None + } + }) + .collect::>(); + + if state_inputs.len() == 1 { + state_inputs.iter().next().map(|&ty| ty.clone()) + } else { + None + } +} + #[test] fn ui() { #[rustversion::stable] diff --git a/axum-macros/src/lib.rs b/axum-macros/src/lib.rs index 2d5983ad..0242ddca 100644 --- a/axum-macros/src/lib.rs +++ b/axum-macros/src/lib.rs @@ -513,12 +513,60 @@ pub fn derive_from_request(item: TokenStream) -> TokenStream { /// async fn handler(request: Request) {} /// ``` /// +/// # Changing state type +/// +/// By default `#[debug_handler]` assumes your state type is `()` unless your handler has a +/// [`axum::extract::State`] argument: +/// +/// ``` +/// use axum::extract::State; +/// # use axum_macros::debug_handler; +/// +/// #[debug_handler] +/// async fn handler( +/// // this makes `#[debug_handler]` use `AppState` +/// State(state): State, +/// ) {} +/// +/// #[derive(Clone)] +/// struct AppState {} +/// ``` +/// +/// If your handler takes multiple [`axum::extract::State`] arguments or you need to otherwise +/// customize the state type you can set it with `#[debug_handler(state = ...)]`: +/// +/// ``` +/// use axum::extract::{State, FromRef}; +/// # use axum_macros::debug_handler; +/// +/// #[debug_handler(state = AppState)] +/// async fn handler( +/// State(app_state): State, +/// State(inner_state): State, +/// ) {} +/// +/// #[derive(Clone)] +/// struct AppState { +/// inner: InnerState, +/// } +/// +/// #[derive(Clone)] +/// struct InnerState {} +/// +/// impl FromRef for InnerState { +/// fn from_ref(state: &AppState) -> Self { +/// state.inner.clone() +/// } +/// } +/// ``` +/// /// # Performance /// /// This macro has no effect when compiled with the release profile. (eg. `cargo build --release`) /// /// [`axum`]: https://docs.rs/axum/latest /// [`Handler`]: https://docs.rs/axum/latest/axum/handler/trait.Handler.html +/// [`axum::extract::State`]: https://docs.rs/axum/0.6/axum/extract/struct.State.html /// [`debug_handler`]: macro@debug_handler #[proc_macro_attribute] pub fn debug_handler(_attr: TokenStream, input: TokenStream) -> TokenStream { diff --git a/axum-macros/tests/debug_handler/fail/duplicate_args.rs b/axum-macros/tests/debug_handler/fail/duplicate_args.rs new file mode 100644 index 00000000..dca335bd --- /dev/null +++ b/axum-macros/tests/debug_handler/fail/duplicate_args.rs @@ -0,0 +1,9 @@ +use axum_macros::debug_handler; + +#[debug_handler(body = BoxBody, body = BoxBody)] +async fn handler() {} + +#[debug_handler(state = (), state = ())] +async fn handler_2() {} + +fn main() {} diff --git a/axum-macros/tests/debug_handler/fail/duplicate_args.stderr b/axum-macros/tests/debug_handler/fail/duplicate_args.stderr new file mode 100644 index 00000000..694b6cb3 --- /dev/null +++ b/axum-macros/tests/debug_handler/fail/duplicate_args.stderr @@ -0,0 +1,11 @@ +error: `body` specified more than once + --> tests/debug_handler/fail/duplicate_args.rs:3:33 + | +3 | #[debug_handler(body = BoxBody, body = BoxBody)] + | ^^^^ + +error: `state` specified more than once + --> tests/debug_handler/fail/duplicate_args.rs:6:29 + | +6 | #[debug_handler(state = (), state = ())] + | ^^^^^ diff --git a/axum-macros/tests/debug_handler/fail/invalid_attrs.stderr b/axum-macros/tests/debug_handler/fail/invalid_attrs.stderr index 20705a0d..93514ebf 100644 --- a/axum-macros/tests/debug_handler/fail/invalid_attrs.stderr +++ b/axum-macros/tests/debug_handler/fail/invalid_attrs.stderr @@ -1,4 +1,4 @@ -error: unknown argument +error: expected `body` or `state` --> tests/debug_handler/fail/invalid_attrs.rs:3:17 | 3 | #[debug_handler(foo)] diff --git a/axum-macros/tests/debug_handler/pass/infer_state.rs b/axum-macros/tests/debug_handler/pass/infer_state.rs new file mode 100644 index 00000000..9f21a8a6 --- /dev/null +++ b/axum-macros/tests/debug_handler/pass/infer_state.rs @@ -0,0 +1,31 @@ +use axum_macros::debug_handler; +use axum::extract::State; + +#[debug_handler] +async fn handler(_: State) {} + +#[debug_handler] +async fn handler_2(_: axum::extract::State) {} + +#[debug_handler] +async fn handler_3( + _: axum::extract::State, + _: axum::extract::State, +) {} + +#[debug_handler] +async fn handler_4( + _: State, + _: State, +) {} + +#[debug_handler] +async fn handler_5( + _: axum::extract::State, + _: State, +) {} + +#[derive(Clone)] +struct AppState; + +fn main() {} diff --git a/axum-macros/tests/debug_handler/pass/set_state.rs b/axum-macros/tests/debug_handler/pass/set_state.rs new file mode 100644 index 00000000..12afaf10 --- /dev/null +++ b/axum-macros/tests/debug_handler/pass/set_state.rs @@ -0,0 +1,27 @@ +use axum_macros::debug_handler; +use axum::extract::{FromRef, FromRequest, RequestParts}; +use axum::async_trait; + +#[debug_handler(state = AppState)] +async fn handler(_: A) {} + +#[derive(Clone)] +struct AppState; + +struct A; + +#[async_trait] +impl FromRequest for A +where + B: Send, + S: Send + Sync, + AppState: FromRef, +{ + type Rejection = (); + + async fn from_request(_req: &mut RequestParts) -> Result { + unimplemented!() + } +} + +fn main() {} diff --git a/axum-macros/tests/debug_handler/pass/state_and_body.rs b/axum-macros/tests/debug_handler/pass/state_and_body.rs new file mode 100644 index 00000000..7e1525f5 --- /dev/null +++ b/axum-macros/tests/debug_handler/pass/state_and_body.rs @@ -0,0 +1,10 @@ +use axum_macros::debug_handler; +use axum::{body::BoxBody, extract::State, http::Request}; + +#[debug_handler(state = AppState, body = BoxBody)] +async fn handler(_: State, _: Request) {} + +#[derive(Clone)] +struct AppState; + +fn main() {}