boitalettres/src/util/stream.rs

55 lines
1.4 KiB
Rust

use std::fmt;
use std::pin::Pin;
use std::task::{self, Poll};
use futures::prelude::*;
use futures::stream::{FuturesOrdered, Stream, StreamFuture};
/// [`SelectAll`](futures::stream::SelectAll) but ordered
#[pin_project::pin_project]
pub struct ConcatAll<St: Stream + Unpin> {
#[pin]
inner: FuturesOrdered<StreamFuture<St>>,
}
impl<St: Stream + Unpin> ConcatAll<St> {
pub fn new() -> Self {
Self {
inner: FuturesOrdered::new(),
}
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn push(&mut self, stream: St) {
use futures::StreamExt;
self.inner.push(stream.into_future());
}
}
impl<St: Stream + Unpin> fmt::Debug for ConcatAll<St> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ConcatAll").finish()
}
}
impl<St: Stream + Unpin> Stream for ConcatAll<St> {
type Item = St::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
loop {
match futures::ready!(self.inner.poll_next_unpin(cx)) {
Some((Some(item), remaining)) => {
self.push(remaining);
return Poll::Ready(Some(item));
}
Some((None, _)) => {}
_ => return Poll::Ready(None),
}
}
}
}