use core::fmt; use core::pin::Pin; use futures_core::future::TryFuture; use futures_core::ready; use futures_core::stream::{FusedStream, Stream, TryStream}; use futures_core::task::{Context, Poll}; #[cfg(feature = "sink")] use futures_sink::Sink; use pin_project_lite::pin_project; pin_project! { /// Stream for the [`try_take_while`](super::TryStreamExt::try_take_while) /// method. #[must_use = "streams do nothing unless polled"] pub struct TryTakeWhile where St: TryStream, { #[pin] stream: St, f: F, #[pin] pending_fut: Option, pending_item: Option, done_taking: bool, } } impl fmt::Debug for TryTakeWhile where St: TryStream + fmt::Debug, St::Ok: fmt::Debug, Fut: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("TryTakeWhile") .field("stream", &self.stream) .field("pending_fut", &self.pending_fut) .field("pending_item", &self.pending_item) .field("done_taking", &self.done_taking) .finish() } } impl TryTakeWhile where St: TryStream, F: FnMut(&St::Ok) -> Fut, Fut: TryFuture, { pub(super) fn new(stream: St, f: F) -> Self { Self { stream, f, pending_fut: None, pending_item: None, done_taking: false, } } delegate_access_inner!(stream, St, ()); } impl Stream for TryTakeWhile where St: TryStream, F: FnMut(&St::Ok) -> Fut, Fut: TryFuture, { type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); if *this.done_taking { return Poll::Ready(None); } Poll::Ready(loop { if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() { let res = ready!(fut.try_poll(cx)); this.pending_fut.set(None); let take = res?; let item = this.pending_item.take(); if take { break item.map(Ok); } else { *this.done_taking = true; break None; } } else if let Some(item) = ready!(this.stream.as_mut().try_poll_next(cx)?) { this.pending_fut.set(Some((this.f)(&item))); *this.pending_item = Some(item); } else { break None; } }) } fn size_hint(&self) -> (usize, Option) { if self.done_taking { return (0, Some(0)); } let pending_len = if self.pending_item.is_some() { 1 } else { 0 }; let (_, upper) = self.stream.size_hint(); let upper = match upper { Some(x) => x.checked_add(pending_len), None => None, }; (0, upper) // can't know a lower bound, due to the predicate } } impl FusedStream for TryTakeWhile where St: TryStream + FusedStream, F: FnMut(&St::Ok) -> Fut, Fut: TryFuture, { fn is_terminated(&self) -> bool { self.done_taking || self.pending_item.is_none() && self.stream.is_terminated() } } // Forwarding impl of Sink from the underlying stream #[cfg(feature = "sink")] impl Sink for TryTakeWhile where S: TryStream + Sink, { type Error = E; delegate_sink!(stream, Item); }