Support using a different rejection for #[derive(FromRequest)] (#1256)

axum-macros/src/from_request/attr.rs

@@ -11,7 +11,11 @@ pub(crate) struct FromRequestFieldAttr {
 }
 
 pub(crate) enum FromRequestContainerAttr {
-    Via(syn::Path),
+    Via {
+        path: syn::Path,
+        rejection: Option<syn::Path>,
+    },
+    Rejection(syn::Path),
     RejectionDerive(kw::rejection_derive, RejectionDeriveOptOuts),
     None,
 }
@@ -19,6 +23,7 @@ pub(crate) enum FromRequestContainerAttr {
 pub(crate) mod kw {
     syn::custom_keyword!(via);
     syn::custom_keyword!(rejection_derive);
+    syn::custom_keyword!(rejection);
     syn::custom_keyword!(Display);
     syn::custom_keyword!(Debug);
     syn::custom_keyword!(Error);
@@ -51,6 +56,7 @@ pub(crate) fn parse_container_attrs(
 
     let mut out_via = None;
     let mut out_rejection_derive = None;
+    let mut out_rejection = None;
 
     // we track the index of the attribute to know which comes last
     // used to give more accurate error messages
@@ -73,11 +79,18 @@ pub(crate) fn parse_container_attrs(
                     out_rejection_derive = Some((idx, rejection_derive, opt_outs));
                 }
             }
+            ContainerAttr::Rejection { rejection, path } => {
+                if out_rejection.is_some() {
+                    return Err(double_attr_error("rejection", rejection));
+                } else {
+                    out_rejection = Some((idx, rejection, path));
+                }
+            }
         }
     }
 
-    match (out_via, out_rejection_derive) {
-        (Some((via_idx, via, _)), Some((rejection_derive_idx, rejection_derive, _))) => {
+    match (out_via, out_rejection_derive, out_rejection) {
+        (Some((via_idx, via, _)), Some((rejection_derive_idx, rejection_derive, _)), _) => {
             if via_idx > rejection_derive_idx {
                 Err(syn::Error::new_spanned(
                     via,
@@ -90,11 +103,41 @@ pub(crate) fn parse_container_attrs(
                 ))
             }
         }
-        (Some((_, _, path)), None) => Ok(FromRequestContainerAttr::Via(path)),
-        (None, Some((_, rejection_derive, opt_outs))) => Ok(
+
+        (
+            _,
+            Some((rejection_derive_idx, rejection_derive, _)),
+            Some((rejection_idx, rejection, _)),
+        ) => {
+            if rejection_idx > rejection_derive_idx {
+                Err(syn::Error::new_spanned(
+                    rejection,
+                    "cannot use both `rejection_derive` and `rejection`",
+                ))
+            } else {
+                Err(syn::Error::new_spanned(
+                    rejection_derive,
+                    "cannot use both `rejection` and `rejection_derive`",
+                ))
+            }
+        }
+
+        (Some((_, _, path)), None, None) => Ok(FromRequestContainerAttr::Via {
+            path,
+            rejection: None,
+        }),
+        (Some((_, _, path)), None, Some((_, _, rejection))) => Ok(FromRequestContainerAttr::Via {
+            path,
+            rejection: Some(rejection),
+        }),
+
+        (None, Some((_, rejection_derive, opt_outs)), _) => Ok(
             FromRequestContainerAttr::RejectionDerive(rejection_derive, opt_outs),
         ),
-        (None, None) => Ok(FromRequestContainerAttr::None),
+
+        (None, None, Some((_, _, rejection))) => Ok(FromRequestContainerAttr::Rejection(rejection)),
+
+        (None, None, None) => Ok(FromRequestContainerAttr::None),
     }
 }
 
@@ -125,6 +168,10 @@ enum ContainerAttr {
         via: kw::via,
         path: syn::Path,
     },
+    Rejection {
+        rejection: kw::rejection,
+        path: syn::Path,
+    },
     RejectionDerive {
         rejection_derive: kw::rejection_derive,
         opt_outs: RejectionDeriveOptOuts,
@@ -147,6 +194,13 @@ impl Parse for ContainerAttr {
                 rejection_derive,
                 opt_outs,
             })
+        } else if lh.peek(kw::rejection) {
+            let rejection = input.parse::<kw::rejection>()?;
+            let content;
+            syn::parenthesized!(content in input);
+            content
+                .parse()
+                .map(|path| Self::Rejection { rejection, path })
         } else {
             Err(lh.error())
         }