atomic_progress/
io.rs

1//! I/O wrappers for tracking data transfer.
2//!
3//! This module provides [`ProgressReader`] and [`ProgressWriter`], which wrap any implementation
4//! of [`std::io::Read`] or [`std::io::Write`].
5//!
6//! # mechanics
7//!
8//! These wrappers act as "pass-through" middleware. Every byte successfully read or written
9//! automatically increments the associated [`Progress`] counter. This is particularly
10//! useful for:
11//!
12//! * File downloads/uploads.
13//! * Hashing large files.
14//! * Compressing/Decompressing data streams.
15//!
16//! The overhead is minimal: a single atomic addition per `read` or `write` call.
17
18use std::io::{self, Read, Write};
19
20use crate::Progress;
21
22/// A wrapper around [`Read`] that increments a [`Progress`] tracker based on bytes read.
23pub struct ProgressReader<R> {
24    inner: R,
25    progress: Progress,
26}
27
28impl<R> ProgressReader<R> {
29    /// Creates a new `ProgressReader` wrapping `inner` with the given `progress` tracker.
30    pub const fn new(inner: R, progress: Progress) -> Self {
31        Self { inner, progress }
32    }
33}
34
35impl<R: Read> Read for ProgressReader<R> {
36    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
37        let n = self.inner.read(buf)?;
38        // cast is safe: we are just tracking bytes
39        self.progress.inc(n as u64);
40        Ok(n)
41    }
42}
43
44/// A wrapper around [`Write`] that increments a [`Progress`] tracker based on bytes written.
45pub struct ProgressWriter<W> {
46    inner: W,
47    progress: Progress,
48}
49
50impl<W> ProgressWriter<W> {
51    /// Creates a new `ProgressWriter` wrapping `inner` with the given `progress` tracker.
52    pub const fn new(inner: W, progress: Progress) -> Self {
53        Self { inner, progress }
54    }
55}
56
57impl<W: Write> Write for ProgressWriter<W> {
58    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
59        let n = self.inner.write(buf)?;
60        self.progress.inc(n as u64);
61        Ok(n)
62    }
63
64    fn flush(&mut self) -> io::Result<()> {
65        self.inner.flush()
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use std::io::{Cursor, Read as _, Write as _};
72
73    use crate::io::{ProgressReader, ProgressWriter};
74
75    use super::Progress;
76
77    /// Reader Tracking
78    /// Verifies bytes read are counted.
79    #[test]
80    fn test_io_reader() {
81        let data = vec![0u8; 100];
82        let p = Progress::new_pb("read", 100u64);
83        let mut reader = ProgressReader::new(Cursor::new(&data), p.clone());
84
85        let mut buf = [0u8; 10];
86        reader.read_exact(&mut buf).unwrap();
87
88        assert_eq!(p.get_pos(), 10);
89    }
90
91    /// Writer Tracking
92    /// Verifies bytes written are counted.
93    #[test]
94    fn test_io_writer() {
95        let p = Progress::new_pb("write", 50u64);
96        let mut writer = ProgressWriter::new(Vec::new(), p.clone());
97
98        writer.write_all(&[1, 2, 3, 4, 5]).unwrap();
99
100        assert_eq!(p.get_pos(), 5);
101    }
102}
atomic_progress/io.rs

atomic_progress/
io.rs
