use core::{ future::Future, marker::PhantomData, pin::Pin, task::{Context, Poll}, }; use pin_project_lite::pin_project; use super::{Service, ServiceFactory}; /// Service for the `map` combinator, changing the type of a service's response. /// /// This is created by the `ServiceExt::map` method. pub struct Map { service: A, f: F, _t: PhantomData<(Req, Res)>, } impl Map { /// Create new `Map` combinator pub(crate) fn new(service: A, f: F) -> Self where A: Service, F: FnMut(A::Response) -> Res, { Self { service, f, _t: PhantomData, } } } impl Clone for Map where A: Clone, F: Clone, { fn clone(&self) -> Self { Map { service: self.service.clone(), f: self.f.clone(), _t: PhantomData, } } } impl Service for Map where A: Service, F: FnMut(A::Response) -> Res + Clone, { type Response = Res; type Error = A::Error; type Future = MapFuture; crate::forward_ready!(service); fn call(&self, req: Req) -> Self::Future { MapFuture::new(self.service.call(req), self.f.clone()) } } pin_project! { pub struct MapFuture where A: Service, F: FnMut(A::Response) -> Res, { f: F, #[pin] fut: A::Future, } } impl MapFuture where A: Service, F: FnMut(A::Response) -> Res, { fn new(fut: A::Future, f: F) -> Self { MapFuture { f, fut } } } impl Future for MapFuture where A: Service, F: FnMut(A::Response) -> Res, { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); match this.fut.poll(cx) { Poll::Ready(Ok(resp)) => Poll::Ready(Ok((this.f)(resp))), Poll::Ready(Err(e)) => Poll::Ready(Err(e)), Poll::Pending => Poll::Pending, } } } /// `MapNewService` new service combinator pub struct MapServiceFactory { a: A, f: F, r: PhantomData<(Res, Req)>, } impl MapServiceFactory { /// Create new `Map` new service instance pub(crate) fn new(a: A, f: F) -> Self where A: ServiceFactory, F: FnMut(A::Response) -> Res, { Self { a, f, r: PhantomData, } } } impl Clone for MapServiceFactory where A: Clone, F: Clone, { fn clone(&self) -> Self { Self { a: self.a.clone(), f: self.f.clone(), r: PhantomData, } } } impl ServiceFactory for MapServiceFactory where A: ServiceFactory, F: FnMut(A::Response) -> Res + Clone, { type Response = Res; type Error = A::Error; type Config = A::Config; type Service = Map; type InitError = A::InitError; type Future = MapServiceFuture; fn new_service(&self, cfg: A::Config) -> Self::Future { MapServiceFuture::new(self.a.new_service(cfg), self.f.clone()) } } pin_project! { pub struct MapServiceFuture where A: ServiceFactory, F: FnMut(A::Response) -> Res, { #[pin] fut: A::Future, f: Option, } } impl MapServiceFuture where A: ServiceFactory, F: FnMut(A::Response) -> Res, { fn new(fut: A::Future, f: F) -> Self { MapServiceFuture { f: Some(f), fut } } } impl Future for MapServiceFuture where A: ServiceFactory, F: FnMut(A::Response) -> Res, { type Output = Result, A::InitError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); if let Poll::Ready(svc) = this.fut.poll(cx)? { Poll::Ready(Ok(Map::new(svc, this.f.take().unwrap()))) } else { Poll::Pending } } } #[cfg(test)] mod tests { use futures_util::future::lazy; use super::*; use crate::{ ok, IntoServiceFactory, Ready, Service, ServiceExt, ServiceFactory, ServiceFactoryExt, }; struct Srv; impl Service<()> for Srv { type Response = (); type Error = (); type Future = Ready>; crate::always_ready!(); fn call(&self, _: ()) -> Self::Future { ok(()) } } #[actix_rt::test] async fn test_poll_ready() { let srv = Srv.map(|_| "ok"); let res = lazy(|cx| srv.poll_ready(cx)).await; assert_eq!(res, Poll::Ready(Ok(()))); } #[actix_rt::test] async fn test_call() { let srv = Srv.map(|_| "ok"); let res = srv.call(()).await; assert!(res.is_ok()); assert_eq!(res.unwrap(), "ok"); } #[actix_rt::test] async fn test_new_service() { let new_srv = (|| ok::<_, ()>(Srv)).into_factory().map(|_| "ok"); let srv = new_srv.new_service(&()).await.unwrap(); let res = srv.call(()).await; assert!(res.is_ok()); assert_eq!(res.unwrap(), ("ok")); } }