http_body_alt/
body_ext.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll, ready},
4};
5
6use pin_project_lite::pin_project;
7
8use super::{body::Body, frame::Frame};
9
10pub trait BodyExt: Body {
11    fn frame(&mut self) -> FrameFuture<'_, Self>
12    where
13        Self: Unpin,
14    {
15        FrameFuture { body: self }
16    }
17
18    fn data(&mut self) -> DataFuture<'_, Self>
19    where
20        Self: Unpin,
21    {
22        DataFuture { body: self }
23    }
24
25    fn chain<B>(self, other: B) -> ChainBody<Self, B>
26    where
27        B: Body<Data = Self::Data>,
28        Self: Sized,
29        Self::Error: From<B::Error>,
30    {
31        ChainBody {
32            first: self,
33            second: other,
34        }
35    }
36}
37
38pub struct FrameFuture<'a, B: ?Sized> {
39    body: &'a mut B,
40}
41
42impl<B> Future for FrameFuture<'_, B>
43where
44    B: Body + Unpin,
45{
46    type Output = Option<Result<Frame<B::Data>, B::Error>>;
47
48    #[inline]
49    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
50        Body::poll_frame(Pin::new(self.get_mut().body), cx)
51    }
52}
53
54pub struct DataFuture<'a, B: ?Sized> {
55    body: &'a mut B,
56}
57
58impl<B> Future for DataFuture<'_, B>
59where
60    B: Body + Unpin,
61{
62    type Output = Option<Result<B::Data, B::Error>>;
63
64    #[inline]
65    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
66        let opt = match ready!(Body::poll_frame(Pin::new(self.get_mut().body), cx)) {
67            Some(Ok(Frame::Data(data))) => Some(Ok(data)),
68            Some(Err(e)) => Some(Err(e)),
69            Some(_) | None => None,
70        };
71
72        Poll::Ready(opt)
73    }
74}
75
76pin_project! {
77    pub struct ChainBody<A, B> {
78        #[pin]
79        first: A,
80        #[pin]
81        second: B
82    }
83}
84
85impl<A, B> Body for ChainBody<A, B>
86where
87    A: Body,
88    B: Body<Data = A::Data>,
89    A::Error: From<B::Error>,
90{
91    type Data = A::Data;
92    type Error = A::Error;
93
94    fn poll_frame(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
95        let this = self.project();
96
97        match ready!(this.first.poll_frame(cx)) {
98            Some(frame) => Poll::Ready(Some(frame)),
99            None => this.second.poll_frame(cx).map_err(Into::into),
100        }
101    }
102
103    #[inline]
104    fn is_end_stream(&self) -> bool {
105        self.first.is_end_stream() && self.second.is_end_stream()
106    }
107
108    fn size_hint(&self) -> crate::SizeHint {
109        self.first.size_hint() + self.second.size_hint()
110    }
111}
112
113impl<B> BodyExt for B where B: Body {}
http_body_alt/body_ext.rs

http_body_alt/
body_ext.rs
