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 20 total, err := c.Reader.Read(d) 21 c.Count += int64(total) 22 c.Err = err 23 24 return total, err 25 } 26 27 // Writer is used to wrap a io.Writer for counting. 28 type Writer struct { 29 io.Writer 30 Count int64 31 Err error 32 } 33 34 // Write implements the io.Writer interface. 35 func (c *Writer) Write(d []byte) (int, error) { 36 if c.Err != nil { 37 return 0, c.Err 38 } 39 40 total, err := c.Writer.Write(d) 41 c.Count += int64(total) 42 c.Err = err 43 44 return total, err 45 } 46 47 // WriteString implements the io.StringWriter interface. 48 func (c *Writer) WriteString(s string) (int, error) { 49 if c.Err != nil { 50 return 0, c.Err 51 } 52 53 total, err := io.WriteString(c.Writer, s) 54 c.Count += int64(total) 55 c.Err = err 56 57 return total, err 58 } 59