1 // Package rwcount implements a simple counter that wraps an io.Reader or io.Writer. 2 // Useful for functions (like binary.Read/Write) which do not return read/write counts. 3 package rwcount // import "vimagination.zapto.org/rwcount" 4 5 import "io" 6 7 // Reader is used to wrap a io.Reader for counting 8 type Reader struct { 9 io.Reader 10 Count int64 11 Err error 12 } 13 14 // Read implements the io.Reader interface 15 func (c *Reader) Read(d []byte) (int, error) { 16 if c.Err != nil { 17 return 0, c.Err 18 } 19 total, err := c.Reader.Read(d) 20 c.Count += int64(total) 21 c.Err = err 22 return total, err 23 } 24 25 // Writer is used to wrap a io.Writer for counting 26 type Writer struct { 27 io.Writer 28 Count int64 29 Err error 30 } 31 32 // Write implements the io.Writer interface 33 func (c *Writer) Write(d []byte) (int, error) { 34 if c.Err != nil { 35 return 0, c.Err 36 } 37 total, err := c.Writer.Write(d) 38 c.Count += int64(total) 39 c.Err = err 40 return total, err 41 } 42 43 // WriteString implements the io.StringWriter interface 44 func (c *Writer) WriteString(s string) (int, error) { 45 if c.Err != nil { 46 return 0, c.Err 47 } 48 total, err := io.WriteString(c.Writer, s) 49 c.Count += int64(total) 50 c.Err = err 51 return total, err 52 }