Skip to main content

rustrade_integration/stream/ext/
mod.rs

1use crate::stream::ext::{forward_by::ForwardBy, forward_clone_by::ForwardCloneBy};
2use futures::Stream;
3
4/// Forward matching items, yield non-matching items.
5pub mod forward_by;
6
7/// Forward clones of matching items, yield all items.
8pub mod forward_clone_by;
9
10/// `Stream` extension trait.
11pub trait BarterStreamExt
12where
13    Self: Stream,
14{
15    /// Add a "consecutive event timeout" to the `Stream`
16    ///
17    /// Upon timeout, the `Stream` ends after executing the provided `TimeoutHandler` function.
18    fn with_timeout<TimeoutHandler>(
19        self,
20        timeout_next_item: std::time::Duration,
21        on_timeout: TimeoutHandler,
22    ) -> impl Stream<Item = Self::Item>
23    where
24        Self: Stream + Sized,
25        TimeoutHandler: FnOnce() + 'static,
26    {
27        use tokio_stream::StreamExt;
28        let mut on_timeout = Some(on_timeout);
29        self.timeout(timeout_next_item)
30            .map_while(move |timeout_result| match timeout_result {
31                Ok(item) => Some(item),
32                Err(_elapsed) => {
33                    if let Some(handler) = on_timeout.take() {
34                        handler();
35                    }
36                    None
37                }
38            })
39    }
40
41    /// Forward a subset of `Stream::Item`s.
42    ///
43    /// All items that the predicate returns as `Left` and forwarded, whilst `Right` items
44    /// continue through the existing `Stream`.
45    fn forward_by<A, B, FnPredicate, FnForward, FwdErr>(
46        self,
47        predicate: FnPredicate,
48        forward: FnForward,
49    ) -> ForwardBy<Self, FnPredicate, FnForward>
50    where
51        Self: Stream + Sized,
52        FnPredicate: Fn(Self::Item) -> futures::future::Either<A, B>,
53        FnForward: FnMut(A) -> Result<(), FwdErr>,
54    {
55        ForwardBy::new(self, predicate, forward)
56    }
57
58    /// Forward a clone of a subset of `Stream::Item`s that match the predicate. The `Stream`
59    /// still yields all items.
60    fn forward_clone_by<FnPredicate, FnForward, FwdErr>(
61        self,
62        predicate: FnPredicate,
63        forward: FnForward,
64    ) -> ForwardCloneBy<Self, FnPredicate, FnForward>
65    where
66        Self: Stream + Sized,
67        Self::Item: Clone,
68        FnPredicate: FnMut(&Self::Item) -> bool,
69        FnForward: FnMut(Self::Item) -> Result<(), FwdErr>,
70    {
71        ForwardCloneBy::new(self, predicate, forward)
72    }
73}
74
75impl<S> BarterStreamExt for S where S: Stream {}