Skip to main content

tuwunel_core/utils/stream/
try_tools.rs

1//! TryStreamTools for futures::TryStream
2#![expect(clippy::type_complexity)]
3
4use futures::{TryStream, TryStreamExt, future, future::Ready, stream::TryTakeWhile};
5
6use crate::Result;
7
8/// Adds general-purpose operations to fallible streams.
9///
10/// The adapters preserve the source error type and successful item order. They
11/// operate lazily without collecting the stream.
12pub trait TryTools<T, E, S>
13where
14	S: TryStream<Ok = T, Error = E, Item = Result<T, E>> + ?Sized,
15	Self: TryStream + Sized,
16{
17	/// Limits the stream to at most `n` successful items.
18	///
19	/// After yielding `n` successes, the adapter consumes one additional
20	/// success to detect the limit. Earlier source errors are still forwarded;
21	/// with zero, they precede consumption of the first unyielded success.
22	fn try_take(
23		self,
24		n: usize,
25	) -> TryTakeWhile<
26		Self,
27		Ready<Result<bool, S::Error>>,
28		impl FnMut(&S::Ok) -> Ready<Result<bool, S::Error>>,
29	>;
30}
31
32impl<T, E, S> TryTools<T, E, S> for S
33where
34	S: TryStream<Ok = T, Error = E, Item = Result<T, E>> + ?Sized,
35	Self: TryStream + Sized,
36{
37	#[inline]
38	fn try_take(
39		self,
40		mut n: usize,
41	) -> TryTakeWhile<
42		Self,
43		Ready<Result<bool, S::Error>>,
44		impl FnMut(&S::Ok) -> Ready<Result<bool, S::Error>>,
45	> {
46		self.try_take_while(move |_| {
47			let res = future::ok(n > 0);
48			n = n.saturating_sub(1);
49			res
50		})
51	}
52}