Skip to content

Commit 64f1ed9

Browse files
committed
Add stream::chain implementation
1 parent 63980d4 commit 64f1ed9

File tree

1 file changed

+29
-0
lines changed

1 file changed

+29
-0
lines changed

src/stream.rs

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,26 @@ pub fn zip<St1, St2>(stream: St1, other: St2) -> impl Stream<Item = (St1::Item,
225225
})
226226
}
227227

228+
pub fn chain<St>(stream: St, other: St) -> impl Stream<Item = St::Item>
229+
where St: Stream,
230+
{
231+
let stream = Box::pin(stream);
232+
let other = Box::pin(other);
233+
let start_with_first = true;
234+
futures::stream::unfold((stream, other, start_with_first), async move | (mut stream, mut other, start_with_first)| {
235+
if start_with_first {
236+
if let Some(item) = await!(next(&mut stream)) {
237+
return Some((item, (stream, other, start_with_first)))
238+
}
239+
}
240+
if let Some(item) = await!(next(&mut other)) {
241+
Some((item, (stream, other, /* start_with_first */ false)))
242+
} else {
243+
None
244+
}
245+
})
246+
}
247+
228248
#[cfg(test)]
229249
mod tests {
230250
use futures::executor;
@@ -374,4 +394,13 @@ mod tests {
374394

375395
assert_eq!(vec![(1, 5), (2, 6), (3, 7)], executor::block_on(collect::<_, Vec<_>>(stream)));
376396
}
397+
398+
#[test]
399+
fn test_chain() {
400+
let stream1 = iter(1..=2);
401+
let stream2 = iter(3..=4);
402+
let stream = chain(stream1, stream2);
403+
404+
assert_eq!(vec![1, 2, 3, 4], executor::block_on(collect::<_, Vec<_>>(stream)));
405+
}
377406
}

0 commit comments

Comments
 (0)