Merge pull request #90 from tamasfe/feat/axum-routes-refactor

feat(aide): sync axum router differences
tamasfe · Nov 30, 2023 · 6eaa058 · 6eaa058
2 parents d05b1fd + 66c88fc
commit 6eaa058
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 33 deletions.
diff --git a/crates/aide/src/axum/mod.rs b/crates/aide/src/axum/mod.rs
@@ -183,7 +183,7 @@ use axum::{
     handler::Handler,
     http::Request,
     response::IntoResponse,
-    routing::{IntoMakeService, Route},
+    routing::{IntoMakeService, Route, RouterAsService, RouterIntoService},
     Router,
 };
 use indexmap::map::Entry;
@@ -458,7 +458,7 @@ where
         self
     }
 
-    /// Alternative to [`nest_service`](Self::nest_service()) which besides nesting the service nests
+    /// Alternative to [`nest_service`](Self::nest_service) which besides nesting the service nests
     /// the generated documentation as well.
     ///
     /// Due to Rust's limitations, currently this function will not
@@ -566,6 +566,22 @@ where
         self.router = self.router.fallback_service(svc);
         self
     }
+
+    /// See [`axum::Router::as_service`] for details.
+    ///
+    /// Using this method will not generate API documentation.
+    #[must_use]
+    pub fn as_service<B>(&mut self) -> RouterAsService<'_, B, S> {
+        self.router.as_service()
+    }
+
+    /// See [`axum::Router::into_service`] for details.
+    ///
+    /// Using this method will not generate API documentation.
+    #[must_use]
+    pub fn into_service<B>(self) -> RouterIntoService<B, S> {
+        self.router.into_service()
+    }
 }
 
 impl ApiRouter<()> {

diff --git a/crates/aide/src/axum/routing.rs b/crates/aide/src/axum/routing.rs
@@ -8,8 +8,8 @@ use crate::{
     openapi::{Operation, PathItem, ReferenceOr, Response, StatusCode},
     Error,
 };
-use axum::body::{Body, HttpBody};
-use axum::routing::Route;
+use axum::routing::{MethodFilter, Route};
+use axum::{body::Body, response::IntoResponse};
 use axum::{
     handler::Handler,
     routing::{self, MethodRouter},
@@ -33,6 +33,15 @@ pub struct ApiMethodRouter<S = (), E = Infallible> {
     pub(crate) router: MethodRouter<S, E>,
 }
 
+impl<S, E> Clone for ApiMethodRouter<S, E> {
+    fn clone(&self) -> Self {
+        Self {
+            operations: self.operations.clone(),
+            router: self.router.clone(),
+        }
+    }
+}
+
 impl<S, E> From<ApiMethodRouter<S, E>> for MethodRouter<S, E> {
     fn from(router: ApiMethodRouter<S, E>) -> Self {
         router.router
@@ -233,48 +242,21 @@ where
     method_router_chain_method!(post, post_with);
     method_router_chain_method!(put, put_with);
     method_router_chain_method!(trace, trace_with);
-
-    /// This method wraps a layer around the [`ApiMethodRouter`]
-    /// For further information see [`axum::routing::method_routing::MethodRouter::layer`]
-    pub fn layer<L, NewError>(self, layer: L) -> ApiMethodRouter<S, NewError>
-    where
-        L: Layer<Route<Infallible>> + Clone + Send + 'static,
-        L::Service: Service<Request<Body>, Response = http::response::Response<Body>, Error = NewError>
-            + Clone
-            + Send
-            + 'static,
-        <L::Service as Service<Request<Body>>>::Future: Send + 'static,
-        Body: HttpBody + 'static,
-        NewError: 'static,
-    {
-        ApiMethodRouter {
-            router: self.router.layer(layer),
-            operations: self.operations,
-        }
-    }
 }
 
 impl<S, E> ApiMethodRouter<S, E>
 where
     S: Clone,
 {
-    /// Create a new, clean [`ApiMethodRouter`] based on [`MethodRouter::new()`](axum::routing::MethodRouter).
+    /// Create a new, empty [`ApiMethodRouter`] based on [`MethodRouter::new()`](axum::routing::MethodRouter).
     pub fn new() -> Self {
         Self {
             operations: IndexMap::default(),
             router: MethodRouter::<S, E>::new(),
         }
     }
-    /// See [`axum::routing::MethodRouter`] and [`axum::extract::State`] for more information.
-    pub fn with_state<S2>(self, state: S) -> ApiMethodRouter<S2, E> {
-        let router = self.router.with_state(state);
-        ApiMethodRouter::<S2, E> {
-            operations: self.operations,
-            router,
-        }
-    }
 
-    /// See [`axum::routing::MethodRouter::merge`] for more information.
+    /// See [`axum::routing::MethodRouter::merge`].
     pub fn merge<M>(mut self, other: M) -> Self
     where
         M: Into<ApiMethodRouter<S, E>>,
@@ -284,6 +266,70 @@ where
         self.router = self.router.merge(other.router);
         self
     }
+
+    /// See [`axum::routing::method_routing::MethodRouter::layer`].
+    pub fn layer<L, NewError>(self, layer: L) -> ApiMethodRouter<S, NewError>
+    where
+        L: Layer<Route<E>> + Clone + Send + 'static,
+        L::Service: Service<Request<Body>> + Clone + Send + 'static,
+        <L::Service as Service<Request<Body>>>::Response: IntoResponse + 'static,
+        <L::Service as Service<Request<Body>>>::Error: Into<NewError> + 'static,
+        <L::Service as Service<Request<Body>>>::Future: Send + 'static,
+        E: 'static,
+        S: 'static,
+        NewError: 'static,
+    {
+        ApiMethodRouter {
+            router: self.router.layer(layer),
+            operations: self.operations,
+        }
+    }
+
+    /// See [`axum::routing::method_routing::MethodRouter::with_state`].
+    pub fn with_state<S2>(self, state: S) -> ApiMethodRouter<S2, E> {
+        ApiMethodRouter {
+            router: self.router.with_state(state),
+            operations: self.operations,
+        }
+    }
+
+    /// See [`axum::routing::method_routing::MethodRouter::on_service`].
+    pub fn on_service<T>(mut self, filter: MethodFilter, svc: T) -> Self
+    where
+        T: Service<Request<Body>, Error = E> + Clone + Send + 'static,
+        T::Response: IntoResponse + 'static,
+        T::Future: Send + 'static,
+    {
+        self.router = self.router.on_service(filter, svc);
+        self
+    }
+
+    /// See [`axum::routing::method_routing::MethodRouter::fallback_service`].
+    pub fn fallback_service<T>(mut self, svc: T) -> Self
+    where
+        T: Service<Request<Body>, Error = E> + Clone + Send + 'static,
+        T::Response: IntoResponse + 'static,
+        T::Future: Send + 'static,
+    {
+        self.router = self.router.fallback_service(svc);
+        self
+    }
+
+    /// See [`axum::routing::method_routing::MethodRouter::route_layer`].
+    pub fn route_layer<L>(self, layer: L) -> ApiMethodRouter<S, E>
+    where
+        L: Layer<Route<E>> + Clone + Send + 'static,
+        L::Service: Service<Request<Body>, Error = E> + Clone + Send + 'static,
+        <L::Service as Service<Request<Body>>>::Response: IntoResponse + 'static,
+        <L::Service as Service<Request<Body>>>::Future: Send + 'static,
+        E: 'static,
+        S: 'static,
+    {
+        ApiMethodRouter {
+            router: self.router.route_layer(layer),
+            operations: self.operations,
+        }
+    }
 }
 
 impl<S, E> Default for ApiMethodRouter<S, E>