1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
use std::{ mem, pin::Pin, task::{Context, Poll}, }; use bytes::Bytes; use futures_core::Stream; use pin_project::pin_project; use crate::error::Error; use super::{Body, BodySize, MessageBody}; #[pin_project(project = ResponseBodyProj)] pub enum ResponseBody<B> { Body(#[pin] B), Other(Body), } impl ResponseBody<Body> { pub fn into_body<B>(self) -> ResponseBody<B> { match self { ResponseBody::Body(b) => ResponseBody::Other(b), ResponseBody::Other(b) => ResponseBody::Other(b), } } } impl<B> ResponseBody<B> { pub fn take_body(&mut self) -> ResponseBody<B> { mem::replace(self, ResponseBody::Other(Body::None)) } } impl<B: MessageBody> ResponseBody<B> { pub fn as_ref(&self) -> Option<&B> { if let ResponseBody::Body(ref b) = self { Some(b) } else { None } } } impl<B: MessageBody> MessageBody for ResponseBody<B> { 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<Option<Result<Bytes, Error>>> { Stream::poll_next(self, cx) } } impl<B: MessageBody> Stream for ResponseBody<B> { type Item = Result<Bytes, Error>; fn poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<Option<Self::Item>> { match self.project() { ResponseBodyProj::Body(body) => body.poll_next(cx), ResponseBodyProj::Other(body) => Pin::new(body).poll_next(cx), } } }