diff --git a/pkg/io/copy/progress/progress_listener.go b/pkg/io/copy/progress/progress_listener.go new file mode 100644 index 0000000..6da6640 --- /dev/null +++ b/pkg/io/copy/progress/progress_listener.go @@ -0,0 +1,37 @@ +package progress + +import "io" + +type WriteProgressCallback func(uint64, uint64) + +type WriteProgressListener struct { + BytesWritten uint64 + ProgressCallback WriteProgressCallback +} + +func (self *WriteProgressListener) Write(p []byte) (int, error) { + n := len(p) + newTotal := self.BytesWritten + uint64(n) + + if self.ProgressCallback != nil { + self.ProgressCallback(self.BytesWritten, newTotal) + } + + self.BytesWritten = newTotal + return n, nil +} + +func (self *WriteProgressListener) WrapReader(reader io.Reader) io.Reader { + return io.TeeReader(reader, self) +} + +func NewProgressListener(callback WriteProgressCallback) *WriteProgressListener { + return &WriteProgressListener{ + BytesWritten: 0, + ProgressCallback: callback, + } +} + +func WrapReader(reader io.Reader, callback WriteProgressCallback) io.Reader { + return NewProgressListener(callback).WrapReader(reader) +}