use std::{ mem, pin::Pin, task::{Context, Poll}, }; use bytes::Bytes; use futures_core::{ready, Stream}; use pin_project::pin_project; use crate::error::Error; use super::{Body, BodySize, MessageBody}; #[pin_project(project = ResponseBodyProj)] pub enum ResponseBody { Body(#[pin] B), Other(Body), } impl ResponseBody { pub fn into_body(self) -> ResponseBody { match self { ResponseBody::Body(b) => ResponseBody::Other(b), ResponseBody::Other(b) => ResponseBody::Other(b), } } } impl ResponseBody { pub fn take_body(&mut self) -> ResponseBody { mem::replace(self, ResponseBody::Other(Body::None)) } } impl ResponseBody { pub fn as_ref(&self) -> Option<&B> { if let ResponseBody::Body(ref b) = self { Some(b) } else { None } } } impl MessageBody for ResponseBody where B: MessageBody, B::Error: Into, { type Error = Error; fn size(&self) -> BodySize { match self { ResponseBody::Body(ref body) => body.size(), ResponseBody::Other(ref body) => body.size(), } } fn poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll>> { Stream::poll_next(self, cx) } } impl Stream for ResponseBody where B: MessageBody, B::Error: Into, { type Item = Result; fn poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { match self.project() { // TODO: MSRV 1.51: poll_map_err ResponseBodyProj::Body(body) => match ready!(body.poll_next(cx)) { Some(Err(err)) => Poll::Ready(Some(Err(err.into()))), Some(Ok(val)) => Poll::Ready(Some(Ok(val))), None => Poll::Ready(None), }, ResponseBodyProj::Other(body) => Pin::new(body).poll_next(cx), } } }