Add clamped reader

1 files changed, 39 insertions, 0 deletions

diff --git a/dataset/clamped.go b/dataset/clamped.go
new file mode 100644
index 0000000..41441b2
--- /dev/null
+++ b/dataset/clamped.go
@@ -0,0 +1,39 @@
+package dataset
+
+import "iter"
+
+type ClampedReader struct {
+	Reader
+	n int
+}
+
+func NewClampedReader(r Reader, n int) *ClampedReader {
+	return &ClampedReader{
+		Reader: r,
+		n:      n,
+	}
+}
+
+func (c *ClampedReader) Num() (int, error) {
+	n, err := c.Reader.Num()
+
+	if err != nil {
+		return 0, err
+	}
+
+	return min(n, c.n), nil
+}
+
+func (c *ClampedReader) Texts() iter.Seq2[int, string] {
+	return func(yield func(int, string) bool) {
+		for n, d := range c.Reader.Texts() {
+			if n >= c.n {
+				return
+			}
+
+			if !yield(n, d) {
+				return
+			}
+		}
+	}
+}