1 | use core::pin::Pin; |
2 | use futures_core::task::{Context, Poll}; |
3 | use futures_core::Stream; |
4 | use pin_project_lite::pin_project; |
5 | |
6 | pin_project! { |
7 | /// Stream for the [poll_immediate](poll_immediate()) function. |
8 | /// |
9 | /// It will never return [Poll::Pending](core::task::Poll::Pending) |
10 | #[derive(Debug, Clone)] |
11 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
12 | pub struct PollImmediate<S> { |
13 | #[pin] |
14 | stream: Option<S> |
15 | } |
16 | } |
17 | |
18 | impl<T, S> Stream for PollImmediate<S> |
19 | where |
20 | S: Stream<Item = T>, |
21 | { |
22 | type Item = Poll<T>; |
23 | |
24 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { |
25 | let mut this = self.project(); |
26 | let stream = match this.stream.as_mut().as_pin_mut() { |
27 | // inner is gone, so we can continue to signal that the stream is closed. |
28 | None => return Poll::Ready(None), |
29 | Some(inner) => inner, |
30 | }; |
31 | |
32 | match stream.poll_next(cx) { |
33 | Poll::Ready(Some(t)) => Poll::Ready(Some(Poll::Ready(t))), |
34 | Poll::Ready(None) => { |
35 | this.stream.set(None); |
36 | Poll::Ready(None) |
37 | } |
38 | Poll::Pending => Poll::Ready(Some(Poll::Pending)), |
39 | } |
40 | } |
41 | |
42 | fn size_hint(&self) -> (usize, Option<usize>) { |
43 | self.stream.as_ref().map_or((0, Some(0)), Stream::size_hint) |
44 | } |
45 | } |
46 | |
47 | impl<S: Stream> super::FusedStream for PollImmediate<S> { |
48 | fn is_terminated(&self) -> bool { |
49 | self.stream.is_none() |
50 | } |
51 | } |
52 | |
53 | /// Creates a new stream that always immediately returns [Poll::Ready](core::task::Poll::Ready) when awaiting it. |
54 | /// |
55 | /// This is useful when immediacy is more important than waiting for the next item to be ready. |
56 | /// |
57 | /// # Examples |
58 | /// |
59 | /// ``` |
60 | /// # futures::executor::block_on(async { |
61 | /// use futures::stream::{self, StreamExt}; |
62 | /// use futures::task::Poll; |
63 | /// |
64 | /// let mut r = stream::poll_immediate(Box::pin(stream::iter(1_u32..3))); |
65 | /// assert_eq!(r.next().await, Some(Poll::Ready(1))); |
66 | /// assert_eq!(r.next().await, Some(Poll::Ready(2))); |
67 | /// assert_eq!(r.next().await, None); |
68 | /// |
69 | /// let mut p = stream::poll_immediate(Box::pin(stream::once(async { |
70 | /// futures::pending!(); |
71 | /// 42_u8 |
72 | /// }))); |
73 | /// assert_eq!(p.next().await, Some(Poll::Pending)); |
74 | /// assert_eq!(p.next().await, Some(Poll::Ready(42))); |
75 | /// assert_eq!(p.next().await, None); |
76 | /// # }); |
77 | /// ``` |
78 | pub fn poll_immediate<S: Stream>(s: S) -> PollImmediate<S> { |
79 | super::assert_stream::<Poll<S::Item>, PollImmediate<S>>(PollImmediate { stream: Some(s) }) |
80 | } |
81 | |