diff --git a/itertools/README.md b/itertools/README.md index b81e673..310258b 100644 --- a/itertools/README.md +++ b/itertools/README.md @@ -2,3 +2,11 @@ A collection of Iterator functions inspired by Python's `itertools` module, adapted to the conventions of Go. + +| Python | Go | Description | +| --- | --- | --- | +| `islice(iterable, stop)` | `Take[T any](n int, seq iter.Seq[T]) iter.Seq[T]` | Yield the first `n` items of the source sequence, then stop. | +| `count(start=0)` | `Count[N Number](start N) iter.Seq[N]` | Yield numbers from `start`, incrementing by 1, without end. | +| `count(start=0, step=1)` | `CountStep[N Number](start, step N) iter.Seq[N]` | Yield numbers from `start`, incrementing by `step`, without end. | +| `repeat(object)` | `RepeatForever[T any](v T) iter.Seq[T]` | Yield the same value over and over, without end. | +| `repeat(object, times)` | `Repeat[T any](v T, times int) iter.Seq[T]` | Yield the same value `times` times, then stop. | diff --git a/itertools/basic.go b/itertools/basic.go index b96c71d..8c21c8c 100644 --- a/itertools/basic.go +++ b/itertools/basic.go @@ -1,12 +1,22 @@ package itertools -import "iter" +import ( + "iter" + + "golang.org/x/exp/constraints" +) + +// Number is the set of types Count and CountStep can generate, i.e. anything +// that supports the + operator with an integer or floating-point step. +type Number interface { + constraints.Integer | constraints.Float +} // Take returns a lazy iterator over the first n values of seq. It stops seq // early once n values have been yielded, so it is safe to apply to infinite -// sequences like Count. Use slices.Collect to materialize the result. -func Take(n int, seq iter.Seq[int]) iter.Seq[int] { - return func(yield func(int) bool) { +// sequences like Count. Use slices.Collect to build a slice with the results. +func Take[T any](n int, seq iter.Seq[T]) iter.Seq[T] { + return func(yield func(T) bool) { if n <= 0 { return } @@ -23,17 +33,38 @@ func Take(n int, seq iter.Seq[int]) iter.Seq[int] { } } +// RepeatForever returns an infinite iterator that yields v forever. +// Use Repeat for the bounded form. +// The caller must break or use a bounded helper (e.g. Take) to stop it. +func RepeatForever[T any](v T) iter.Seq[T] { + return func(yield func(T) bool) { + for yield(v) { + } + } +} + +// Repeat returns an iterator that yields v exactly times times. +// A non-positive times yields nothing. +func Repeat[T any](v T, times int) iter.Seq[T] { + return func(yield func(T) bool) { + for range times { + if !yield(v) { + return + } + } + } +} + // Count returns an infinite iterator that yields start, start+1, start+2, … -// It is equivalent to Python's itertools.count(start). -func Count(start int) iter.Seq[int] { +func Count[N Number](start N) iter.Seq[N] { return CountStep(start, 1) } // CountStep returns an infinite iterator that yields start, start+step, -// start+2*step, … — the Go equivalent of Python's itertools.count(start, step). +// start+2*step, … // The caller must break or use a bounded helper (e.g. Take) to stop it. -func CountStep(start, step int) iter.Seq[int] { - return func(yield func(int) bool) { +func CountStep[N Number](start, step N) iter.Seq[N] { + return func(yield func(N) bool) { for i := start; ; i += step { if !yield(i) { return diff --git a/itertools/basic_test.go b/itertools/basic_test.go index 101c7e1..86cb39e 100644 --- a/itertools/basic_test.go +++ b/itertools/basic_test.go @@ -13,6 +13,47 @@ func TestCount(t *testing.T) { } } +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 diff --git a/itertools/go.mod b/itertools/go.mod index 8b54b0e..7ef38cc 100644 --- a/itertools/go.mod +++ b/itertools/go.mod @@ -1,3 +1,5 @@ module ramalho.org/itertools go 1.26.2 + +require golang.org/x/exp v0.0.0-20260603202125-055de637280b