count_write/
lib.rs

1// count-write
2// Copyright (C) SOFe
3//
4// Licensed under the Apache License, Version 2.0 (the License);
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an AS IS BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use std::io::Write;
17
18type Result<T = (), E = std::io::Error> = std::result::Result<T, E>;
19
20/// A wrapper for [`std::io::Write`][Write] (and other optional Write implementations) that counts
21/// the total number of bytes written successfully.
22///
23/// ```
24/// use std::io::{BufWriter, Write};
25///
26/// use count_write::CountWrite;
27///
28/// let vec = vec![]; // note that Vec<u8> : Write
29/// let cw = CountWrite::from(vec);
30/// let mut bw = BufWriter::new(cw);
31/// write!(bw, "abc")?;
32/// assert_eq!(bw.into_inner()?.count(), 3);
33/// # Ok::<(), std::io::Error>(())
34/// ```
35///
36/// If the inner writer does not successfully write all bytes,
37/// only the successfully written bytes are counted.
38/// (This does not affect users who always use [`write_all`][write_all], [`write!`][write-macro],
39/// etc.)
40///
41/// ```
42/// use std::io::{Result, Write};
43///
44/// use count_write::CountWrite;;
45///
46/// /// A dummy struct that only accepts half of the input into a Vec.
47/// struct OnlyHalf;
48///
49/// impl Write for OnlyHalf {
50///     fn write(&mut self, buf: &[u8]) -> Result<usize> {
51///     dbg!(buf);
52///         Ok((buf.len() + 1) / 2)
53///     }
54///
55///     fn flush(&mut self) -> Result<()> { Ok(()) }
56/// }
57///
58/// let mut cw = CountWrite::from(OnlyHalf);
59/// write!(cw, "abc")?; // Here, we keep retrying writing into the writer
60/// assert_eq!(cw.count(), 3);
61///
62/// let mut cw = CountWrite::from(OnlyHalf);
63/// cw.write(b"abc")?; // Here, we only write to the writer once and do not retry
64/// assert_eq!(cw.count(), 2);
65/// # Ok::<(), std::io::Error>(())
66/// ```
67///
68/// [Write]: https://doc.rust-lang.org/std/io/trait.Write.html
69/// [write_all]: https://doc.rust-lang.org/std/io/trait.Write.html#method.write_all
70/// [write-macro]: https://doc.rust-lang.org/std/macro.write.html
71pub struct CountWrite<W> {
72    inner: W,
73    count: u64,
74}
75
76impl<W> CountWrite<W> {
77    /// Returns the number of bytes successfull written so far
78    pub fn count(&self) -> u64 {
79        self.count
80    }
81
82    /// Extracts the inner writer, discarding this wrapper
83    pub fn into_inner(self) -> W {
84        self.inner
85    }
86}
87
88impl<W> From<W> for CountWrite<W> {
89    fn from(inner: W) -> Self {
90        Self { inner, count: 0 }
91    }
92}
93
94impl<W: Write> Write for CountWrite<W> {
95    fn write(&mut self, buf: &[u8]) -> Result<usize> {
96        let written = self.inner.write(buf)?;
97        self.count += written as u64;
98        Ok(written)
99    }
100
101    fn flush(&mut self) -> Result {
102        self.inner.flush()
103    }
104}
105
106#[cfg(feature = "futures")]
107mod futures;
108
109#[cfg(feature = "tokio")]
110mod tokio;
count_write/lib.rs

count_write/
lib.rs
