diff --git a/stream/stream.go b/stream/stream.go index c694aaa..eb18ad0 100644 --- a/stream/stream.go +++ b/stream/stream.go @@ -167,7 +167,7 @@ func (s stream[T]) Peek(consumer func(item T)) stream[T] { return s } -// Skip Returns a stream consisting of the remaining elements of this stream after discarding the first n elements of the stream. If this stream contains fewer than n elements then an empty stream will be returned. +// Skip returns a stream consisting of the remaining elements of this stream after discarding the first n elements of the stream. If this stream contains fewer than n elements then an empty stream will be returned. func (s stream[T]) Skip(n int) stream[T] { if n <= 0 { return s @@ -180,8 +180,6 @@ func (s stream[T]) Skip(n int) stream[T] { return FromSlice(source) } - // source = make([]T, l-n+1, l-n+1) - for i := n; i < l; i++ { source = append(source, s.source[i]) } @@ -189,6 +187,25 @@ func (s stream[T]) Skip(n int) stream[T] { return FromSlice(source) } +// Limit returns a stream consisting of the elements of this stream, truncated to be no longer than maxSize in length. +func (s stream[T]) Limit(maxSize int) stream[T] { + if s.source == nil { + return s + } + + if maxSize < 0 { + return FromSlice([]T{}) + } + + source := make([]T, 0, maxSize) + + for i := 0; i < len(s.source) && i < maxSize; i++ { + source = append(source, s.source[i]) + } + + return FromSlice(source) +} + // Count returns the count of elements in the stream. func (s stream[T]) Count() int { return len(s.source) diff --git a/stream/stream_test.go b/stream/stream_test.go index 4fb5e38..7adc716 100644 --- a/stream/stream_test.go +++ b/stream/stream_test.go @@ -151,8 +151,23 @@ func TestStream_Skip(t *testing.T) { s1 := stream.Skip(-1) s2 := stream.Skip(0) - // s2 := stream.Skip(0) assert.Equal([]int{1, 2, 3, 4, 5, 6}, s1.ToSlice()) assert.Equal([]int{1, 2, 3, 4, 5, 6}, s2.ToSlice()) } + +func TestStream_Limit(t *testing.T) { + assert := internal.NewAssert(t, "TestStream_Limit") + + stream := FromSlice([]int{1, 2, 3, 4, 5, 6}) + + s1 := stream.Limit(-1) + s2 := stream.Limit(0) + s3 := stream.Limit(1) + s4 := stream.Limit(6) + + assert.Equal([]int{}, s1.ToSlice()) + assert.Equal([]int{}, s2.ToSlice()) + assert.Equal([]int{1}, s3.ToSlice()) + assert.Equal([]int{1, 2, 3, 4, 5, 6}, s4.ToSlice()) +}