Skip to main content

async_iter_ext/iter/
process_result.rs

1use std::{
2    fmt::{Debug, Formatter},
3    ops::Deref,
4    pin::{Pin, pin},
5    task::{Context, Poll},
6};
7
8use crate::AsyncIterator;
9
10/// Defines the strategy to use when processing results from an asynchronous iterator.
11#[derive(Default, Clone, Debug)]
12pub enum ProcessResultsStrategy {
13    /// Continue processing all results, separating successes and errors.
14    #[default]
15    Partition,
16
17    /// Stop processing at the first error encountered.
18    BreakOnError,
19}
20
21/// A container that holds both successful and erroneous results.
22pub struct ProcessResultsContainer<T, E> {
23    successes: Vec<T>,
24    errors: Vec<E>,
25}
26
27impl<T, E> Deref for ProcessResultsContainer<T, E> {
28    type Target = Vec<T>;
29
30    /// Dereferences to the vector of successful results.
31    fn deref(&self) -> &Self::Target {
32        self.successes()
33    }
34}
35
36impl<T, E> ProcessResultsContainer<T, E> {
37    /// Converts the container into a `Result`, returning `Ok` with successes,
38    /// or the first `Err` if there are any errors.
39    pub fn into_result(self) -> Result<Vec<T>, E> {
40        if !self.errors.is_empty() {
41            Err(self.into_errors().remove(0))
42        } else {
43            Ok(self.successes)
44        }
45    }
46
47    /// Consumes the container and returns the vector of successes.
48    pub fn into_successes(self) -> Vec<T> {
49        self.successes
50    }
51
52    /// Consumes the container and returns the vector of errors.
53    pub fn into_errors(self) -> Vec<E> {
54        self.errors
55    }
56
57    /// Returns a reference to the vector of successes.
58    pub fn successes(&self) -> &Vec<T> {
59        self.successes.as_ref()
60    }
61
62    /// Returns a reference to the vector of errors.
63    pub fn errors(&self) -> &Vec<E> {
64        self.errors.as_ref()
65    }
66}
67
68impl<T, E> Debug for ProcessResultsContainer<T, E>
69where
70    T: Debug,
71    E: Debug,
72{
73    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
74        f.debug_struct("ProcessResultsContainer")
75            .field("successes", &self.successes)
76            .field("errors", &self.errors)
77            .finish()
78    }
79}
80
81impl<T, E> Clone for ProcessResultsContainer<T, E>
82where
83    T: Clone,
84    E: Clone,
85{
86    fn clone(&self) -> Self {
87        Self {
88            successes: self.successes.clone(),
89            errors: self.errors.clone(),
90        }
91    }
92}
93
94impl<T, E> From<(Vec<T>, Vec<E>)> for ProcessResultsContainer<T, E> {
95    /// Creates a `ProcessResultsContainer` from a tuple of successes and errors.
96    fn from((successes, errors): (Vec<T>, Vec<E>)) -> Self {
97        Self { successes, errors }
98    }
99}
100
101/// A future that processes results from an asynchronous iterator,
102/// collecting successes and errors based on the specified strategy.
103pub struct ProcessResults<I, T, E>
104where
105    I: AsyncIterator<Item = Result<T, E>>,
106{
107    iter: I,
108    strategy: ProcessResultsStrategy,
109}
110
111impl<I, T, E> ProcessResults<I, T, E>
112where
113    I: AsyncIterator<Item = Result<T, E>>,
114{
115    /// Constructs a new `ProcessResults` future using the given async iterator.
116    pub fn new(iter: I) -> ProcessResults<I, T, E> {
117        Self {
118            iter,
119            strategy: ProcessResultsStrategy::default(),
120        }
121    }
122
123    /// Sets the processing strategy to use for handling errors during iteration.
124    pub fn with_process_strategy(mut self, strategy: ProcessResultsStrategy) -> Self {
125        self.strategy = strategy;
126        self
127    }
128}
129
130impl<I, T, E> Future for ProcessResults<I, T, E>
131where
132    I: AsyncIterator<Item = Result<T, E>> + Unpin,
133    T: Unpin,
134    E: Unpin,
135{
136    type Output = ProcessResultsContainer<T, E>;
137
138    /// Polls the future and returns a container of results.
139    /// Depending on the strategy, it either:
140    /// - `Partition`: Collects all successes and errors.
141    /// - `BreakOnError`: Stops at the first error and returns it immediately.
142    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
143        let strategy = self.strategy.clone();
144        let mut pinned_fut = pin!(self.get_mut().iter.sync_iter());
145
146        loop {
147            match pinned_fut.as_mut().poll(cx) {
148                Poll::Pending => {}
149                Poll::Ready(res) => {
150                    let mut successes = vec![];
151                    let mut errors = vec![];
152
153                    for item in res {
154                        match item {
155                            Ok(item) => successes.push(item),
156                            Err(error) => {
157                                errors.push(error);
158                                match strategy {
159                                    ProcessResultsStrategy::Partition => {}
160                                    ProcessResultsStrategy::BreakOnError => {
161                                        return Poll::Ready((vec![], errors).into());
162                                    }
163                                }
164                            }
165                        }
166                    }
167
168                    return Poll::Ready((successes, errors).into());
169                }
170            }
171        }
172    }
173}
174
175impl<I, T, E> Debug for ProcessResults<I, T, E>
176where
177    I: AsyncIterator<Item = Result<T, E>> + Debug,
178{
179    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
180        f.debug_struct("ProcessResults")
181            .field("iter", &self.iter)
182            .field("strategy", &self.strategy)
183            .finish()
184    }
185}