stream: add stream extend function

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<E, S>(buff: &mut E, b: S)
+where
+    S: Stream,
+    E: Extend<S::Item>,
+{
+    crate::pin!(b);
+    while let Some(item) = b.next().await {
+        buff.extend(std::iter::once(item))
+    }
+}

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};
 

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<String, i32> = vec![("one".to_string(), 1)].into_iter().collect();
+    let s = stream::iter(hm);
+
+    let mut buff: HashMap<String, i32> = 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);
+}