atomic_progress/
iter.rs

1//! Iterator adapters for automatic progress tracking.
2//!
3//! This module provides the [`ProgressIteratorExt`] trait, which adds helper methods
4//! to any Rust [`Iterator`]. This allows you to attach a progress bar to a loop with
5//! a single method call.
6//!
7//! # Heuristics
8//!
9//! The adapters automatically check [`Iterator::size_hint`]:
10//! * If the iterator provides an exact upper bound, a **Bar** is created with that total.
11//! * If the bounds are unknown (or `(0, None)`), a **Spinner** is created.
12//!
13//! # Example
14//!
15//! ```ignore
16//! use atomic_progress::ProgressIteratorExt;
17//!
18//! // Automatically becomes a Bar because vec.len() is known
19//! for item in vec![1, 2, 3].into_iter().progress() {
20//!     // ...
21//! }
22//! ```
23
24use compact_str::CompactString;
25
26use crate::{
27    progress::{Progress, ProgressType},
28    stack::ProgressStack,
29};
30
31/// An iterator adapter that wraps an underlying iterator and tracks progress.
32///
33/// Increments the progress position on every call to `next()`.
34pub struct ProgressIter<I> {
35    iter: I,
36    progress: Progress,
37}
38
39impl<I> ProgressIter<I> {
40    /// Creates a new `ProgressIter`.
41    ///
42    /// Note: This is usually constructed via [`ProgressIteratorExt`] methods.
43    pub const fn new(iter: I, progress: Progress) -> Self {
44        Self { iter, progress }
45    }
46}
47
48impl<I: Iterator> Iterator for ProgressIter<I> {
49    type Item = I::Item;
50
51    fn next(&mut self) -> Option<Self::Item> {
52        let item = self.iter.next();
53
54        if item.is_some() {
55            // Immediate atomic update for strict accuracy
56            self.progress.inc(1u64);
57        } else {
58            // Iterator exhausted
59            self.progress.finish();
60        }
61
62        item
63    }
64}
65
66/// Extension trait to easily attach progress tracking to any Iterator.
67pub trait ProgressIteratorExt: Sized {
68    /// Wraps the iterator in a [`Progress`].
69    ///
70    /// Automatically detects if the iterator has a known size (via `size_hint`) to
71    /// choose between a Bar or Spinner.
72    fn progress(self) -> ProgressIter<Self>;
73
74    /// Wraps the iterator in a [`Progress`] with a specific name.
75    fn progress_with_name(self, name: impl Into<CompactString>) -> ProgressIter<Self>;
76
77    /// Wraps the iterator using an existing [`Progress`] instance.
78    fn progress_with(self, progress: Progress) -> ProgressIter<Self>;
79
80    /// Creates a new [`Progress`] attached to the given [`ProgressStack`] and wraps the iterator.
81    fn progress_in(self, stack: &ProgressStack) -> ProgressIter<Self>;
82
83    /// Internal helper to determine progress type and total from `size_hint`.
84    fn type_from_size_hint(&self) -> (ProgressType, u64);
85}
86
87impl<I: Iterator> ProgressIteratorExt for I {
88    fn progress(self) -> ProgressIter<Self> {
89        self.progress_with_name(CompactString::default())
90    }
91
92    fn progress_with_name(self, name: impl Into<CompactString>) -> ProgressIter<Self> {
93        let (kind, total) = self.type_from_size_hint();
94        let progress = Progress::new(kind, name, total);
95        ProgressIter::new(self, progress)
96    }
97
98    fn progress_with(self, progress: Progress) -> ProgressIter<Self> {
99        ProgressIter::new(self, progress)
100    }
101
102    fn progress_in(self, stack: &ProgressStack) -> ProgressIter<Self> {
103        let (kind, total) = self.type_from_size_hint();
104        // Stack methods return the Arc<Progress>, which we pass to the iterator
105        let progress = match kind {
106            ProgressType::Bar => stack.add_pb(CompactString::default(), total),
107            ProgressType::Spinner => stack.add_spinner(CompactString::default()),
108        };
109        ProgressIter::new(self, progress)
110    }
111
112    fn type_from_size_hint(&self) -> (ProgressType, u64) {
113        let (lower, upper) = self.size_hint();
114        // If the upper bound matches the lower bound, we have an exact size -> Bar.
115        // Otherwise -> Spinner.
116        match upper {
117            Some(u) if u == lower => (ProgressType::Bar, u as u64),
118            _ => (ProgressType::Spinner, 0),
119        }
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::ProgressIteratorExt as _;
126
127    /// Iterator Integration
128    /// Verifies the extension trait correctly wraps and tracks an iterator.
129    #[test]
130    fn test_iterator_adapter() {
131        let data = [1, 2, 3, 4, 5];
132        let mut count = 0;
133
134        // Create iterator, wrap in progress, consume it
135        let iter = data.iter().progress_with_name("iter_test");
136        let progress_handle = iter.progress.clone(); // Clone handle to inspect state
137
138        for _ in iter {
139            count += 1;
140        }
141
142        assert_eq!(count, 5);
143        assert_eq!(progress_handle.get_pos(), 5);
144        assert!(
145            progress_handle.is_finished(),
146            "Iterator exhaustion should finish progress"
147        );
148        assert_eq!(
149            progress_handle.get_total(),
150            5,
151            "Total should be inferred from Vec len"
152        );
153    }
154}
atomic_progress/iter.rs

atomic_progress/
iter.rs
