1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
use crate::Stream;
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::task::{Context, Poll};
pin_project! {
/// A `Stream` that wraps the values in an `Option`.
///
/// Whenever the wrapped stream yields an item, this stream yields that item
/// wrapped in `Some`. When the inner stream ends, then this stream first
/// yields a `None` item, and then this stream will also end.
///
/// # Example
///
/// Using `StreamNotifyClose` to handle closed streams with `StreamMap`.
///
/// ```
/// use tokio_stream::{StreamExt, StreamMap, StreamNotifyClose};
///
/// #[tokio::main]
/// async fn main() {
/// let mut map = StreamMap::new();
/// let stream = StreamNotifyClose::new(tokio_stream::iter(vec![0, 1]));
/// let stream2 = StreamNotifyClose::new(tokio_stream::iter(vec![0, 1]));
/// map.insert(0, stream);
/// map.insert(1, stream2);
/// while let Some((key, val)) = map.next().await {
/// match val {
/// Some(val) => println!("got {val:?} from stream {key:?}"),
/// None => println!("stream {key:?} closed"),
/// }
/// }
/// }
/// ```
#[must_use = "streams do nothing unless polled"]
pub struct StreamNotifyClose<S> {
#[pin]
inner: Option<S>,
}
}
impl<S> StreamNotifyClose<S> {
/// Create a new `StreamNotifyClose`.
pub fn new(stream: S) -> Self {
Self {
inner: Some(stream),
}
}
/// Get back the inner `Stream`.
///
/// Returns `None` if the stream has reached its end.
pub fn into_inner(self) -> Option<S> {
self.inner
}
}
impl<S> Stream for StreamNotifyClose<S>
where
S: Stream,
{
type Item = Option<S::Item>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
// We can't invoke poll_next after it ended, so we unset the inner stream as a marker.
match self
.as_mut()
.project()
.inner
.as_pin_mut()
.map(|stream| S::poll_next(stream, cx))
{
Some(Poll::Ready(Some(item))) => Poll::Ready(Some(Some(item))),
Some(Poll::Ready(None)) => {
self.project().inner.set(None);
Poll::Ready(Some(None))
}
Some(Poll::Pending) => Poll::Pending,
None => Poll::Ready(None),
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
if let Some(inner) = &self.inner {
// We always return +1 because when there's stream there's atleast one more item.
let (l, u) = inner.size_hint();
(l.saturating_add(1), u.and_then(|u| u.checked_add(1)))
} else {
(0, Some(0))
}
}
}