#21 Add Take implements

filter.go

@@ -11,7 +11,7 @@ type filterIterator[T any] struct {
 	current    T
 }
 
-// Filter makes Iterable contains elements which filterFunc is true.
+// Filter makes Iterable with elements which filterFunc is true.
 //
 //   itb := gcf.FromSlice([]int{1, 2, 3})
 //   itb = gcf.Filter(itb, func(v int) bool { return v%2 > 0 })

take.go

@@ -0,0 +1,57 @@
+package gcf
+
+type takeIterable[T any] struct {
+	itb   Iterable[T]
+	count int
+}
+
+type takeIterator[T any] struct {
+	it      Iterator[T]
+	count   int
+	i       int
+	done    bool
+	current T
+}
+
+// Take makes Iterable with count elements from ahead.
+//
+//   itb := gcf.FromSlice([]{1, 2, 3})
+//   itb = gcf.Take(itb, 2)
+//
+// If count is 0 or negative, returns empty Iterable.
+func Take[T any](itb Iterable[T], count int) Iterable[T] {
+	if itb == nil {
+		return empty[T]()
+	}
+	if count < 1 {
+		return empty[T]()
+	}
+	return &takeIterable[T]{itb, count}
+}
+
+func (itb *takeIterable[T]) Iterator() Iterator[T] {
+	return &takeIterator[T]{itb.itb.Iterator(), itb.count, 0, false, zero[T]()}
+}
+
+func (it *takeIterator[T]) MoveNext() bool {
+	if it.done {
+		return false
+	}
+	if it.count <= it.i {
+		it.done = true
+		it.current = zero[T]()
+		return false
+	}
+	it.i++
+	if it.it.MoveNext() {
+		it.current = it.it.Current()
+		return true
+	}
+	it.done = true
+	it.current = zero[T]()
+	return false
+}
+
+func (it *takeIterator[T]) Current() T {
+	return it.current
+}

take_test.go

@@ -0,0 +1,97 @@
+package gcf_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/meian/gcf"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTake(t *testing.T) {
+	type args struct {
+		itb   gcf.Iterable[int]
+		count int
+	}
+	tests := []struct {
+		name string
+		args args
+		want []int
+	}{
+		{
+			name: "take 1 from slice 3",
+			args: args{
+				itb:   gcf.FromSlice([]int{1, 2, 3}),
+				count: 1,
+			},
+			want: []int{1},
+		},
+		{
+			name: "take 2 from slice 3",
+			args: args{
+				itb:   gcf.FromSlice([]int{1, 2, 3}),
+				count: 2,
+			},
+			want: []int{1, 2},
+		},
+		{
+			name: "take 3 from slice 3",
+			args: args{
+				itb:   gcf.FromSlice([]int{1, 2, 3}),
+				count: 3,
+			},
+			want: []int{1, 2, 3},
+		},
+		{
+			name: "take 4 from slice 3",
+			args: args{
+				itb:   gcf.FromSlice([]int{1, 2, 3}),
+				count: 4,
+			},
+			want: []int{1, 2, 3},
+		},
+		{
+			name: "take 0 from slice 3",
+			args: args{
+				itb:   gcf.FromSlice([]int{1, 2, 3}),
+				count: 0,
+			},
+			want: []int{},
+		},
+		{
+			name: "take negative",
+			args: args{
+				itb:   gcf.FromSlice([]int{1, 2, 3}),
+				count: -1,
+			},
+			want: []int{},
+		},
+		{
+			name: "nil Iterable",
+			args: args{
+				itb:   nil,
+				count: 3,
+			},
+			want: []int{},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			itb := gcf.Take(tt.args.itb, tt.args.count)
+			s := gcf.ToSlice(itb)
+			assert.Equal(t, tt.want, s)
+		})
+	}
+
+	itb := gcf.FromSlice([]int{1, 2, 3})
+	itb = gcf.Take(itb, 2)
+	testBeforeAndAfter(t, itb)
+}
+
+func ExampleTake() {
+	itb := gcf.FromSlice([]int{1, 2, 3, 4, 5})
+	itb = gcf.Take(itb, 3)
+	fmt.Println(gcf.ToSlice(itb))
+	// Output:
+	// [1 2 3]
+}