diff --git a/stream/stream.go b/stream/stream.go index d2f789e..f02fd03 100644 --- a/stream/stream.go +++ b/stream/stream.go @@ -2,7 +2,7 @@ // Use of this source code is governed by MIT license // Package stream implements a sequence of elements supporting sequential and parallel aggregate operations. -// this package is a experiment to explore if stream in go can work as the way java does. it's complete, but not +// this package is an experiment to explore if stream in go can work as the way java does. it's complete, but not // powerful like other libs package stream @@ -51,6 +51,11 @@ type stream[T any] struct { source []T } +// Of creates a stream stream whose elements are the specified values. +func Of[T any](elems ...T) stream[T] { + return FromSlice(elems) +} + // FromSlice create stream from slice. func FromSlice[T any](source []T) stream[T] { return stream[T]{source: source} @@ -102,6 +107,35 @@ func hashKey(data any) string { return buffer.String() } +// Filter returns a stream consisting of the elements of this stream that match the given predicate. +func (s stream[T]) Filter(predicate func(item T) bool) stream[T] { + source := make([]T, 0) + + for _, v := range s.source { + if predicate(v) { + source = append(source, v) + } + } + + return FromSlice(source) +} + +// Map returns a stream consisting of the elements of this stream that apply the given function to elements of stream. +func (s stream[T]) Map(mapper func(item T) T) stream[T] { + source := make([]T, s.Count(), s.Count()) + + for i, v := range s.source { + source[i] = mapper(v) + } + + return FromSlice(source) +} + +// Count returns the count of elements in the stream. +func (s stream[T]) Count() int { + return len(s.source) +} + // ToSlice return the elements in the stream. func (s stream[T]) ToSlice() []T { return s.source diff --git a/stream/stream_test.go b/stream/stream_test.go index 3b91dbd..948ccbb 100644 --- a/stream/stream_test.go +++ b/stream/stream_test.go @@ -55,3 +55,31 @@ func TestStream_Distinct(t *testing.T) { // {[{001 Tom 10} {002 Jim 20} {003 Mike 30}]} t.Log(distinctStream) } + +func TestStream_Filter(t *testing.T) { + assert := internal.NewAssert(t, "TestStream_Filter") + + stream := FromSlice([]int{1, 2, 3, 4, 5}) + + isEven := func(n int) bool { + return n%2 == 0 + } + + even := stream.Filter(isEven) + + assert.Equal([]int{2, 4}, even.ToSlice()) +} + +func TestStream_Map(t *testing.T) { + assert := internal.NewAssert(t, "TestStream_Map") + + stream := FromSlice([]int{1, 2, 3}) + + addOne := func(n int) int { + return n + 1 + } + + s := stream.Map(addOne) + + assert.Equal([]int{2, 3, 4}, s.ToSlice()) +}