diff options
| -rw-r--r-- | dataset/clamped.go | 39 |
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 + } + } + } +} |
