62 lines
1.5 KiB
Go
62 lines
1.5 KiB
Go
package countwriter
|
|
|
|
import (
|
|
"io"
|
|
)
|
|
|
|
// CountWriter wraps an io.Writer, see [NewCountWriter].
|
|
//
|
|
// It is intended to be used to count different writes to an underlying writer.
|
|
// Once an error occurs, no more writes are passed through, and the underlying error is returned instead.
|
|
// This means that in practice, calls to write can be continued and are ignored silently.
|
|
//
|
|
// The underlying sum of bytes written and error can be seen using [Sum].
|
|
type CountWriter struct {
|
|
w io.Writer
|
|
|
|
n int
|
|
err error
|
|
}
|
|
|
|
// NewCountWriter creates a new [CountWriter] that delegates to w.
|
|
func NewCountWriter(w io.Writer) *CountWriter {
|
|
return &CountWriter{w: w}
|
|
}
|
|
|
|
// write performs the write operation w on this writer.
|
|
func (cw *CountWriter) write(w func() (int, error)) (int, error) {
|
|
// if there was an error, return it and don't do a write
|
|
if cw.err != nil {
|
|
return 0, cw.err
|
|
}
|
|
|
|
// call the writer
|
|
n, err := w()
|
|
|
|
// update the underling state
|
|
cw.n += n
|
|
cw.err = err
|
|
|
|
// and return
|
|
return n, err
|
|
}
|
|
|
|
// Write implements [io.Writer]
|
|
func (cw *CountWriter) Write(p []byte) (int, error) {
|
|
return cw.write(func() (int, error) {
|
|
return cw.w.Write(p)
|
|
})
|
|
}
|
|
|
|
// WriteString implements [io.WriteString].
|
|
// See [Write].
|
|
func (cw *CountWriter) WriteString(s string) (int, error) {
|
|
return cw.write(func() (int, error) {
|
|
return io.WriteString(cw.w, s)
|
|
})
|
|
}
|
|
|
|
// Sum returns the state, that is the total number of bytes written and any error
|
|
func (cw *CountWriter) Sum() (int, error) {
|
|
return cw.n, cw.err
|
|
}
|