diff options
author | Artem Vorotnikov <artem@vorotnikov.me> | 2019-12-24 19:20:02 +0300 |
---|---|---|
committer | Carl Lerche <me@carllerche.com> | 2019-12-24 08:20:02 -0800 |
commit | 101f770af33ae65820e1cc0e9b89d998b3c1317a (patch) | |
tree | 98d3137a8be58b451ba29e7120ce0afe4517fcc8 /tokio/src/stream | |
parent | 6ff4e349e28a4d89098f2587e70c86281c2ae182 (diff) |
stream: add StreamExt::take (#2025)
Diffstat (limited to 'tokio/src/stream')
-rw-r--r-- | tokio/src/stream/mod.rs | 30 | ||||
-rw-r--r-- | tokio/src/stream/take.rs | 76 |
2 files changed, 106 insertions, 0 deletions
diff --git a/tokio/src/stream/mod.rs b/tokio/src/stream/mod.rs index eecc7a7c..f29791b4 100644 --- a/tokio/src/stream/mod.rs +++ b/tokio/src/stream/mod.rs @@ -22,6 +22,9 @@ use next::Next; mod try_next; use try_next::TryNext; +mod take; +use take::Take; + pub use futures_core::Stream; /// An extension trait for `Stream`s that provides a variety of convenient @@ -202,6 +205,33 @@ pub trait StreamExt: Stream { { FilterMap::new(self, f) } + + /// Creates a new stream of at most `n` items of the underlying stream. + /// + /// Once `n` items have been yielded from this stream then it will always + /// return that the stream is done. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use tokio::stream::{self, StreamExt}; + /// + /// let mut stream = stream::iter(1..=10).take(3); + /// + /// assert_eq!(Some(1), stream.next().await); + /// assert_eq!(Some(2), stream.next().await); + /// assert_eq!(Some(3), stream.next().await); + /// assert_eq!(None, stream.next().await); + /// # } + /// ``` + fn take(self, n: usize) -> Take<Self> + where + Self: Sized, + { + Take::new(self, n) + } } impl<T: ?Sized> StreamExt for T where T: Stream {} diff --git a/tokio/src/stream/take.rs b/tokio/src/stream/take.rs new file mode 100644 index 00000000..f0dbbb00 --- /dev/null +++ b/tokio/src/stream/take.rs @@ -0,0 +1,76 @@ +use crate::stream::Stream; + +use core::cmp; +use core::fmt; +use core::pin::Pin; +use core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream for the [`map`](super::StreamExt::map) method. + #[must_use = "streams do nothing unless polled"] + pub struct Take<St> { + #[pin] + stream: St, + remaining: usize, + } +} + +impl<St> fmt::Debug for Take<St> +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Take") + .field("stream", &self.stream) + .finish() + } +} + +impl<St: Stream> Take<St> { + pub(super) fn new(stream: St, remaining: usize) -> Self { + Self { stream, remaining } + } +} + +impl<St> Stream for Take<St> +where + St: Stream, +{ + type Item = St::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + if *self.as_mut().project().remaining > 0 { + self.as_mut().project().stream.poll_next(cx).map(|ready| { + match &ready { + Some(_) => { + *self.as_mut().project().remaining -= 1; + } + None => { + *self.as_mut().project().remaining = 0; + } + } + ready + }) + } else { + Poll::Ready(None) + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.remaining == 0 { + return (0, Some(0)); + } + + let (lower, upper) = self.stream.size_hint(); + + let lower = cmp::min(lower, self.remaining as usize); + + let upper = match upper { + Some(x) if x < self.remaining as usize => Some(x), + _ => Some(self.remaining as usize), + }; + + (lower, upper) + } +} |