From 3e89644297dbba121434c4877748b4485fbca1cc Mon Sep 17 00:00:00 2001 From: Sumner Evans Date: Thu, 30 May 2024 13:16:15 -0600 Subject: [PATCH 1/2] progress: add Reader and Writer Signed-off-by: Sumner Evans --- go.mod | 2 +- progress/doc.go | 11 +++++++++++ progress/reader.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++ progress/writer.go | 43 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 progress/doc.go create mode 100644 progress/reader.go create mode 100644 progress/writer.go diff --git a/go.mod b/go.mod index b324859..c3fc127 100644 --- a/go.mod +++ b/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 ) diff --git a/progress/doc.go b/progress/doc.go new file mode 100644 index 0000000..b7160f3 --- /dev/null +++ b/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 diff --git a/progress/reader.go b/progress/reader.go new file mode 100644 index 0000000..a26cf72 --- /dev/null +++ b/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) diff --git a/progress/writer.go b/progress/writer.go new file mode 100644 index 0000000..e731be4 --- /dev/null +++ b/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) From 54fbd25e269c19dfaa522e3a84e0d925d07a810b Mon Sep 17 00:00:00 2001 From: Sumner Evans Date: Thu, 30 May 2024 13:39:28 -0600 Subject: [PATCH 2/2] progress: add tests Signed-off-by: Sumner Evans --- progress/reader_test.go | 26 ++++++++++++++++++++++++++ progress/writer_test.go | 26 ++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) create mode 100644 progress/reader_test.go create mode 100644 progress/writer_test.go diff --git a/progress/reader_test.go b/progress/reader_test.go new file mode 100644 index 0000000..bcdeedd --- /dev/null +++ b/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) +} diff --git a/progress/writer_test.go b/progress/writer_test.go new file mode 100644 index 0000000..03bb06d --- /dev/null +++ b/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) +}