azure_speech/stream_ext.rs
1use crate::callback::Callback;
2use core::fmt;
3use core::pin::Pin;
4use core::task::{Context, Poll};
5use pin_project_lite::pin_project;
6use std::future::Future;
7use std::pin::pin;
8use tokio_stream::{Stream, StreamExt as _};
9
10pin_project! {
11/// Stream for the [`stop_after`](stop_after) method.
12#[must_use = "streams do nothing unless polled"]
13 pub struct StopAfter<St, F> {
14 #[pin]
15 stream: St,
16 predicate: F,
17 done: bool,
18 }
19}
20
21impl<St, F> StopAfter<St, F> {
22 pub(super) fn new(stream: St, predicate: F) -> Self {
23 Self {
24 stream,
25 predicate,
26 done: false,
27 }
28 }
29}
30
31impl<St, F> fmt::Debug for StopAfter<St, F>
32where
33 St: fmt::Debug,
34{
35 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36 f.debug_struct("StopAfter")
37 .field("stream", &self.stream)
38 .field("done", &self.done)
39 .finish()
40 }
41}
42
43impl<St, F> Stream for StopAfter<St, F>
44where
45 St: Stream,
46 F: FnMut(&St::Item) -> bool,
47{
48 type Item = St::Item;
49
50 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
51 if !*self.as_mut().project().done {
52 self.as_mut().project().stream.poll_next(cx).map(|ready| {
53 let ready = ready.map(|item| {
54 if (self.as_mut().project().predicate)(&item) {
55 *self.as_mut().project().done = true;
56 }
57 item
58 });
59 ready
60 })
61 } else {
62 Poll::Ready(None)
63 }
64 }
65
66 fn size_hint(&self) -> (usize, Option<usize>) {
67 if self.done {
68 return (0, Some(0));
69 }
70
71 let (_, upper) = self.stream.size_hint();
72
73 (0, upper)
74 }
75}
76
77/// An extension trait for `Stream` that provides a variety of convenient combinator functions.
78pub trait StreamExt: Stream
79where
80 Self: 'static,
81{
82 /// Takes elements from this stream until the provided predicate resolves to `true`.
83 ///
84 /// This function operates similarly to `Iterator::take_while`, extracting elements from the
85 /// stream until the predicate `f` evaluates to `false`. Unlike `Iterator::take_while`, this function
86 /// also returns the last evaluated element for which the predicate was `true`, marking the stream as done afterwards.
87 /// Once an element causes the predicate to return false, the stream will consistently return that it is finished.
88 ///
89 /// # Examples
90 ///
91 /// Basic usage:
92 ///
93 /// ```
94 /// use tokio_stream::{self as stream, StreamExt as _};
95 /// use azure_speech::StreamExt;
96 ///
97 /// #[tokio::main]
98 /// async fn main() {
99 ///
100 /// let mut stream = stream::iter(1..=5).stop_after(|&x| x >= 3);
101 ///
102 /// assert_eq!(Some(1), stream.next().await);
103 /// assert_eq!(Some(2), stream.next().await);
104 /// assert_eq!(Some(3), stream.next().await);
105 /// // Since 4 > 3, the stream is now considered done
106 /// assert_eq!(None, stream.next().await);
107 /// }
108 /// ```
109 ///
110 /// This function is particularly useful when you need to process elements of a stream up to a certain point,
111 /// and then stop processing, including the element that caused the stop condition.
112 fn stop_after<F>(self, f: F) -> StopAfter<Self, F>
113 where
114 F: FnMut(&Self::Item) -> bool,
115 Self: Sized,
116 {
117 StopAfter::new(self, f)
118 }
119
120 /// Calls the provided callback for each item in the stream.
121 fn use_callbacks<C>(self, callback: C) -> impl Future<Output = ()>
122 where
123 Self: Sized + Send + Sync,
124 C: Callback<Item = Self::Item> + 'static,
125 {
126 async move {
127 let mut _self = pin!(self);
128 while let Some(event) = _self.next().await {
129 callback.on_event(event).await;
130 }
131 }
132 }
133}
134
135impl<St: ?Sized + 'static> StreamExt for St where St: Stream {}