Skip to main content

tokio_stream/stream_ext/
filter.rs

1use crate::Stream;
2
3use core::fmt;
4use core::pin::Pin;
5use core::task::{ready, Context, Poll};
6use futures_core::FusedStream;
7use pin_project_lite::pin_project;
8
9pin_project! {
10    /// Stream returned by the [`filter`](super::StreamExt::filter) method.
11    #[must_use = "streams do nothing unless polled"]
12    pub struct Filter<St, F> {
13        #[pin]
14        stream: St,
15        f: F,
16    }
17}
18
19impl<St, F> fmt::Debug for Filter<St, F>
20where
21    St: fmt::Debug,
22{
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        f.debug_struct("Filter")
25            .field("stream", &self.stream)
26            .finish()
27    }
28}
29
30impl<St, F> Filter<St, F> {
31    pub(super) fn new(stream: St, f: F) -> Self {
32        Self { stream, f }
33    }
34}
35
36impl<St, F> Stream for Filter<St, F>
37where
38    St: Stream,
39    F: FnMut(&St::Item) -> bool,
40{
41    type Item = St::Item;
42
43    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> {
44        loop {
45            match ready!(self.as_mut().project().stream.poll_next(cx)) {
46                Some(e) => {
47                    if (self.as_mut().project().f)(&e) {
48                        return Poll::Ready(Some(e));
49                    }
50                }
51                None => return Poll::Ready(None),
52            }
53        }
54    }
55
56    fn size_hint(&self) -> (usize, Option<usize>) {
57        (0, self.stream.size_hint().1) // can't know a lower bound, due to the predicate
58    }
59}
60
61impl<St, F> FusedStream for Filter<St, F>
62where
63    St: FusedStream,
64    F: FnMut(&St::Item) -> bool,
65{
66    fn is_terminated(&self) -> bool {
67        self.stream.is_terminated()
68    }
69}