progress: add Reader and Writer

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/rs/zerolog v1.32.0
 	github.com/stretchr/testify v1.9.0
 	golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842
+	golang.org/x/sys v0.12.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -16,5 +17,4 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	golang.org/x/sys v0.12.0 // indirect
 )

progress/doc.go

@@ -0,0 +1,11 @@
+// Copyright (c) 2024 Sumner Evans
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+// Package progress provides wrappers for [io.Writer] and [io.Reader] that
+// report the progress of the read or write operation via a callback.
+package progress
+
+const defaultUpdateInterval = 256 * 1024

progress/reader.go

@@ -0,0 +1,47 @@
+package progress
+
+import "io"
+
+// Reader is an [io.Reader] that reports the number of bytes read from it via a
+// callback. The callback is called at most every "updateInterval" bytes. The
+// updateInterval can be set using the [Reader.WithUpdateInterval] method.
+//
+// The following is an example of how to use [Reader] to report the progress of
+// reading from a file:
+//
+//	file, _ := os.Open("file.txt")
+//	progressReader := NewReader(f, func(readBytes int) {
+//		fmt.Printf("Read %d bytes\n", readBytes)
+//	})
+//	io.ReadAll(progressReader)
+type Reader struct {
+	inner          io.Reader
+	readBytes      int
+	progressFn     func(readBytes int)
+	lastUpdate     int
+	updateInterval int
+}
+
+func NewReader(r io.Reader, progressFn func(readBytes int)) *Reader {
+	return &Reader{inner: r, progressFn: progressFn, updateInterval: defaultUpdateInterval}
+}
+
+func (r *Reader) WithUpdateInterval(bytes int) *Reader {
+	r.updateInterval = bytes
+	return r
+}
+
+func (r *Reader) Read(p []byte) (n int, err error) {
+	n, err = r.inner.Read(p)
+	if err != nil {
+		return n, err
+	}
+	r.readBytes += n
+	if r.lastUpdate == 0 || r.readBytes-r.lastUpdate > r.updateInterval {
+		r.progressFn(r.readBytes)
+		r.lastUpdate = r.readBytes
+	}
+	return n, nil
+}
+
+var _ io.Reader = (*Reader)(nil)

progress/reader_test.go

@@ -0,0 +1,26 @@
+package progress_test
+
+import (
+	"bytes"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.mau.fi/util/progress"
+)
+
+func TestReader(t *testing.T) {
+	reader := bytes.NewReader(bytes.Repeat([]byte{42}, 1024*1024))
+
+	var progressUpdates []int
+	progressReader := progress.NewReader(reader, func(readBytes int) {
+		progressUpdates = append(progressUpdates, readBytes)
+	})
+
+	data, err := io.ReadAll(progressReader)
+	assert.NoError(t, err)
+	assert.Equal(t, data, bytes.Repeat([]byte{42}, 1024*1024))
+
+	assert.Greater(t, len(progressUpdates), 1)
+	assert.IsIncreasing(t, progressUpdates)
+}

progress/writer.go

@@ -0,0 +1,43 @@
+package progress
+
+import "io"
+
+// Writer is an [io.Writer] that reports the number of bytes written to it via
+// a callback. The callback is called at most every "updateInterval" bytes. The
+// updateInterval can be set using the [Writer.WithUpdateInterval] method.
+//
+// The following is an example of how to use [Writer] to report the progress of
+// writing to a file:
+//
+//	file, _ := os.Create("file.txt")
+//	progressWriter := progress.NewWriter(func(processedBytes int) {
+//	    fmt.Printf("Processed %d bytes\n", processedBytes)
+//	})
+//	writerWithProgress := io.MultiWriter(file, progressWriter)
+//	io.Copy(writerWithProgress, bytes.NewReader(bytes.Repeat([]byte{42}, 1024*1024)))
+type Writer struct {
+	processedBytes int
+	progressFn     func(processedBytes int)
+	lastUpdate     int
+	updateInterval int
+}
+
+func NewWriter(progressFn func(processedBytes int)) *Writer {
+	return &Writer{progressFn: progressFn, updateInterval: defaultUpdateInterval}
+}
+
+func (w *Writer) WithUpdateInterval(bytes int) *Writer {
+	w.updateInterval = bytes
+	return w
+}
+
+func (w *Writer) Write(p []byte) (n int, err error) {
+	w.processedBytes += len(p)
+	if w.lastUpdate == 0 || w.processedBytes-w.lastUpdate > w.updateInterval {
+		w.progressFn(w.processedBytes)
+		w.lastUpdate = w.processedBytes
+	}
+	return len(p), nil
+}
+
+var _ io.Writer = (*Writer)(nil)

progress/writer_test.go

@@ -0,0 +1,26 @@
+package progress_test
+
+import (
+	"bytes"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.mau.fi/util/progress"
+)
+
+func TestWriter(t *testing.T) {
+	var progressUpdates []int
+	progressWriter := progress.NewWriter(func(processedBytes int) {
+		progressUpdates = append(progressUpdates, processedBytes)
+	})
+
+	for i := 0; i < 10; i++ {
+		_, err := io.Copy(progressWriter, bytes.NewReader(bytes.Repeat([]byte{42}, 256*1024)))
+		require.NoError(t, err)
+	}
+
+	assert.Greater(t, len(progressUpdates), 1)
+	assert.IsIncreasing(t, progressUpdates)
+}