Something went wrong. Try again.
Showcasing the power of iterators
Something went wrong. Try again.
3.6 kB · 149 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150package itertools
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 items of seq.// It stops seq early once n items have been yielded,// so it is safe to apply to infinite 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) { taken := 0 for v := range seq { if taken >= n { return } if !yield(v) { return } taken++ } }}
// 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, …// The caller must break or use a bounded helper (e.g. Take) to stop it.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 caller must break or use a bounded helper (e.g. Take) to stop it.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 } } }}
// Cycle returns an infinite iterator that yields the items of seq forever.// Iterates over seq once, saving a copy of each item to yield each later.// Memory grows with the length of seq; an infinite seq makes the saved slice// grow without bound and never reaches the replay loop.// An empty seq yields nothing.// The caller must break or use a bounded helper (e.g. Take) to stop it.func Cycle[T any](seq iter.Seq[T]) iter.Seq[T] { return func(yield func(T) bool) { var saved []T for v := range seq { if !yield(v) { return } saved = append(saved, v) } if len(saved) == 0 { return } for { for _, v := range saved { if !yield(v) { return } } } }}
// Chain returns an iterator that yields the items of each seq in turn,// exhausting the first before moving to the second, and so on.// With no arguments it yields nothing.func Chain[T any](seqs ...iter.Seq[T]) iter.Seq[T] { return func(yield func(T) bool) { for _, seq := range seqs { for v := range seq { if !yield(v) { return } } } }}
// Filter returns an iterator that yields only the items of seq for which// predicate returns true.func Filter[T any](predicate func(T) bool, seq iter.Seq[T]) iter.Seq[T] { return func(yield func(T) bool) { for v := range seq { if predicate(v) { if !yield(v) { return } } } }}
// DropWhile discards items from seq while the predicate returns true.// Then it yields all remaining items, regardless of the predicate.func DropWhile[T any](predicate func(T) bool, seq iter.Seq[T]) iter.Seq[T] { return func(yield func(T) bool) { dropping := true for v := range seq { if dropping { if predicate(v) { continue } dropping = false } if !yield(v) { return } } }}