use std::io;
pub struct CountingWriter<W> {
wtr: W,
cnt: u64,
}
impl<W: io::Write> CountingWriter<W> {
pub fn new(wtr: W) -> CountingWriter<W> {
CountingWriter {
wtr: wtr,
cnt: 0,
}
}
pub fn count(&self) -> u64 {
self.cnt
}
pub fn into_inner(self) -> W {
self.wtr
}
}
impl<W: io::Write> io::Write for CountingWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let n = try!(self.wtr.write(buf));
self.cnt += n as u64;
Ok(n)
}
fn flush(&mut self) -> io::Result<()> {
self.wtr.flush()
}
}
#[cfg(test)]
mod tests {
use std::io::Write;
use super::CountingWriter;
#[test]
fn counts_bytes() {
let mut wtr = CountingWriter::new(vec![]);
wtr.write_all(b"foobar").unwrap();
assert_eq!(wtr.count(), 6);
}
}