diff --git a/tokio/src/stream/extend.rs b/tokio/src/stream/extend.rs new file mode 100644 index 00000000000..d398b56de93 --- /dev/null +++ b/tokio/src/stream/extend.rs @@ -0,0 +1,22 @@ +use crate::stream::{Stream, StreamExt}; +use std::iter::Extend; +/// Extends a slice from a Stream +/// # Example +/// `rust,no_run` +/// use tokio::stream; +/// +/// let buff = vec![-2]; +/// let s = stream::iter(vec![0, 2, 4, 6]); +/// +/// stream::extend(&mut buff, s).await; +/// assert_eq!(vec![-2, 0, 2, 4, 6], buff); +pub async fn extend(buff: &mut E, b: S) +where + S: Stream, + E: Extend, +{ + crate::pin!(b); + while let Some(item) = b.next().await { + buff.extend(std::iter::once(item)) + } +} diff --git a/tokio/src/stream/mod.rs b/tokio/src/stream/mod.rs index 6bf4232494a..dbd6bcf782b 100644 --- a/tokio/src/stream/mod.rs +++ b/tokio/src/stream/mod.rs @@ -66,6 +66,8 @@ mod collect; use collect::Collect; pub use collect::FromStream; +mod extend; +pub use extend::extend; mod empty; pub use empty::{empty, Empty}; diff --git a/tokio/tests/stream_extend.rs b/tokio/tests/stream_extend.rs new file mode 100644 index 00000000000..2be1732c441 --- /dev/null +++ b/tokio/tests/stream_extend.rs @@ -0,0 +1,26 @@ +use std::collections::HashMap; + +use tokio::stream; + +#[tokio::test] +async fn vec_extend_stream() { + let s = stream::iter(vec![0, 2, 4, 6]); + let mut buff = vec![-2]; + stream::extend(&mut buff, s).await; + assert_eq!(buff, vec![-2, 0, 2, 4, 6]); +} + +#[tokio::test] +async fn hash_map_extend_stream() { + let hm: HashMap = vec![("one".to_string(), 1)].into_iter().collect(); + let s = stream::iter(hm); + + let mut buff: HashMap = vec![("zero".to_string(), 0)].into_iter().collect(); + buff.insert("zero".to_string(), 0); + stream::extend(&mut buff, s).await; + + let hm_expected: HashMap<_, _> = vec![("zero".to_string(), 0), ("one".to_string(), 1)] + .into_iter() + .collect(); + assert_eq!(hm_expected, buff); +}