Something went wrong. Try again.
Showcasing the power of iterators
Something went wrong. Try again.
5.7 kB · 234 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235package itertools
import ( "iter" "slices" "testing")
func TestCount(t *testing.T) { got := slices.Collect(Take(5, Count(3))) want := []int{3, 4, 5, 6, 7} if !slices.Equal(got, want) { t.Errorf("got %v, want %v", got, want) }}
func TestCountFloat(t *testing.T) { // A float start makes N=float64, and the untyped step 1 inside Count is // inferred as float64 too, so the sequence advances by 1.0. got := slices.Collect(Take(4, Count(1.5))) want := []float64{1.5, 2.5, 3.5, 4.5} if !slices.Equal(got, want) { t.Errorf("got %v, want %v", got, want) }}
func TestRepeat(t *testing.T) { cases := []struct { name string v string times int want []string }{ {"three times", "Ni!", 3, []string{"Ni!", "Ni!", "Ni!"}}, {"once", "x", 1, []string{"x"}}, {"zero", "x", 0, nil}, {"negative", "x", -2, nil}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { got := slices.Collect(Repeat(tc.v, tc.times)) if !slices.Equal(got, tc.want) { t.Errorf("got %v, want %v", got, tc.want) } }) }}
func TestRepeatForever(t *testing.T) { // RepeatForever is infinite, so bound it with Take. got := slices.Collect(Take(4, RepeatForever(7))) want := []int{7, 7, 7, 7} if !slices.Equal(got, want) { t.Errorf("got %v, want %v", got, want) }}
func TestCountStep(t *testing.T) { cases := []struct { name string start int step int len int want []int }{ {"step 1", 0, 1, 5, []int{0, 1, 2, 3, 4}}, {"non-zero start", 10, 1, 4, []int{10, 11, 12, 13}}, {"step > 1", 0, 10, 5, []int{0, 10, 20, 30, 40}}, {"negative step", 10, -3, 4, []int{10, 7, 4, 1}}, {"negative start", -5, 2, 4, []int{-5, -3, -1, 1}}, {"take zero", 0, 1, 0, nil}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { got := slices.Collect(Take(tc.len, CountStep(tc.start, tc.step))) if !slices.Equal(got, tc.want) { t.Errorf("got %v, want %v", got, tc.want) } }) }}
func TestCountStep_earlyBreak(t *testing.T) { // Verify the iterator stops cleanly when the caller breaks. count := 0 for v := range CountStep(1, 1) { count++ if v == 5 { break } } if count != 5 { t.Errorf("expected 5 iterations, got %d", count) }}
func TestCycle(t *testing.T) { cases := []struct { name string in []int len int want []int }{ {"repeats", []int{1, 2, 3}, 7, []int{1, 2, 3, 1, 2, 3, 1}}, {"single value", []int{9}, 4, []int{9, 9, 9, 9}}, {"shorter than input", []int{1, 2, 3}, 2, []int{1, 2}}, {"exact length", []int{1, 2, 3}, 3, []int{1, 2, 3}}, {"empty input", []int{}, 5, nil}, {"take zero", []int{1, 2, 3}, 0, nil}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { got := slices.Collect(Take(tc.len, Cycle(slices.Values(tc.in)))) if !slices.Equal(got, tc.want) { t.Errorf("got %v, want %v", got, tc.want) } }) }}
func TestCycle_earlyBreak(t *testing.T) { // Verify the iterator stops cleanly when the caller breaks mid-cycle. count := 0 for v := range Cycle(slices.Values([]int{1, 2, 3})) { count++ if v == 2 && count > 3 { break } } if count != 5 { t.Errorf("expected 5 iterations, got %d", count) }}
func TestChain(t *testing.T) { cases := []struct { name string in [][]int want []int }{ {"two sequences", [][]int{{1, 2}, {3, 4}}, []int{1, 2, 3, 4}}, {"three sequences", [][]int{{1}, {2, 3}, {4, 5, 6}}, []int{1, 2, 3, 4, 5, 6}}, {"one sequence", [][]int{{1, 2, 3}}, []int{1, 2, 3}}, {"no sequences", nil, nil}, {"empty in the middle", [][]int{{1, 2}, {}, {3}}, []int{1, 2, 3}}, {"all empty", [][]int{{}, {}}, nil}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { seqs := make([]iter.Seq[int], 0, len(tc.in)) for _, s := range tc.in { seqs = append(seqs, slices.Values(s)) } got := slices.Collect(Chain(seqs...)) if !slices.Equal(got, tc.want) { t.Errorf("got %v, want %v", got, tc.want) } }) }}
func TestChain_earlyBreak(t *testing.T) { // Verify Chain stops the current seq and skips later seqs when the caller // breaks. The second sequence is infinite, so reaching it without stopping // would hang the test. count := 0 for v := range Chain(slices.Values([]int{1, 2, 3}), Count(100)) { count++ if v == 2 { break } } if count != 2 { t.Errorf("expected 2 iterations, got %d", count) }}
func TestFilter(t *testing.T) { isEven := func(n int) bool { return n%2 == 0 } cases := []struct { name string in []int want []int }{ {"evens", []int{1, 2, 3, 4, 5, 6}, []int{2, 4, 6}}, {"none match", []int{1, 3, 5}, nil}, {"all match", []int{2, 4, 6}, []int{2, 4, 6}}, {"empty input", []int{}, nil}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { got := slices.Collect(Filter(isEven, slices.Values(tc.in))) if !slices.Equal(got, tc.want) { t.Errorf("got %v, want %v", got, tc.want) } }) }}
func TestFilter_earlyBreak(t *testing.T) { isEven := func(n int) bool { return n%2 == 0 } count := 0 for range Filter(isEven, Count(1)) { count++ if count == 3 { break } } if count != 3 { t.Errorf("expected 3 iterations, got %d", count) }}
func TestDropWhile(t *testing.T) { small := func(n int) bool { return n < 10 } cases := []struct { name string in []int want []int }{ {"drop two", []int{1, 2, 10, 11}, []int{10, 11}}, {"ignore predicate", []int{1, 10, 5}, []int{10, 5}}, {"drop all", []int{2, 4, 6}, nil}, {"empty input", []int{}, nil}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { got := slices.Collect(DropWhile(small, slices.Values(tc.in))) if !slices.Equal(got, tc.want) { t.Errorf("got %v, want %v", got, tc.want) } }) }}