mirror of
https://github.com/duke-git/lancet.git
synced 2026-02-04 12:52:28 +08:00
feat: add take iterator
This commit is contained in:
@@ -112,3 +112,28 @@ func Reduce[T any, U any](iter Iterator[T], initial U, reducer func(U, T) U) U {
|
|||||||
|
|
||||||
return acc
|
return acc
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func Take[T any](it Iterator[T], num int) Iterator[T] {
|
||||||
|
return &takeIterator[T]{it: it, num: num}
|
||||||
|
}
|
||||||
|
|
||||||
|
type takeIterator[T any] struct {
|
||||||
|
it Iterator[T]
|
||||||
|
num int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (iter *takeIterator[T]) Next() (T, bool) {
|
||||||
|
if iter.num <= 0 {
|
||||||
|
var zero T
|
||||||
|
return zero, false
|
||||||
|
}
|
||||||
|
item, ok := iter.it.Next()
|
||||||
|
if ok {
|
||||||
|
iter.num--
|
||||||
|
}
|
||||||
|
return item, ok
|
||||||
|
}
|
||||||
|
|
||||||
|
func (iter *takeIterator[T]) HasNext() bool {
|
||||||
|
return iter.num > 0
|
||||||
|
}
|
||||||
|
|||||||
@@ -60,3 +60,14 @@ func TestReduce(t *testing.T) {
|
|||||||
sum := Reduce(iter, 0, func(a, b int) int { return a + b })
|
sum := Reduce(iter, 0, func(a, b int) int { return a + b })
|
||||||
assert.Equal(10, sum)
|
assert.Equal(10, sum)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestTakeIterator(t *testing.T) {
|
||||||
|
assert := internal.NewAssert(t, "TestTakeIterator")
|
||||||
|
|
||||||
|
iter := FromSlice([]int{1, 2, 3, 4, 5})
|
||||||
|
|
||||||
|
iter = Take(iter, 3)
|
||||||
|
|
||||||
|
result := ToSlice(iter)
|
||||||
|
assert.Equal([]int{1, 2, 3}, result)
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user