tokio_stream/stream_ext/
take.rs

1use crate::Stream;
2
3use core::cmp;
4use core::fmt;
5use core::pin::Pin;
6use core::task::{Context, Poll};
7use futures_core::FusedStream;
8use pin_project_lite::pin_project;
9
10pin_project! {
11    /// Stream for the [`take`](super::StreamExt::take) method.
12    #[must_use = "streams do nothing unless polled"]
13    pub struct Take<St> {
14        #[pin]
15        stream: St,
16        remaining: usize,
17    }
18}
19
20impl<St> fmt::Debug for Take<St>
21where
22    St: fmt::Debug,
23{
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        f.debug_struct("Take")
26            .field("stream", &self.stream)
27            .finish()
28    }
29}
30
31impl<St> Take<St> {
32    pub(super) fn new(stream: St, remaining: usize) -> Self {
33        Self { stream, remaining }
34    }
35}
36
37impl<St> Stream for Take<St>
38where
39    St: Stream,
40{
41    type Item = St::Item;
42
43    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
44        if *self.as_mut().project().remaining > 0 {
45            self.as_mut().project().stream.poll_next(cx).map(|ready| {
46                match &ready {
47                    Some(_) => {
48                        *self.as_mut().project().remaining -= 1;
49                    }
50                    None => {
51                        *self.as_mut().project().remaining = 0;
52                    }
53                }
54                ready
55            })
56        } else {
57            Poll::Ready(None)
58        }
59    }
60
61    fn size_hint(&self) -> (usize, Option<usize>) {
62        if self.remaining == 0 {
63            return (0, Some(0));
64        }
65
66        let (lower, upper) = self.stream.size_hint();
67
68        let lower = cmp::min(lower, self.remaining);
69
70        let upper = match upper {
71            Some(x) if x < self.remaining => Some(x),
72            _ => Some(self.remaining),
73        };
74
75        (lower, upper)
76    }
77}
78
79impl<St> FusedStream for Take<St>
80where
81    St: Stream,
82{
83    fn is_terminated(&self) -> bool {
84        self.remaining == 0
85    }
86}
tokio_stream/stream_ext/take.rs

tokio_stream/stream_ext/
take.rs
