mediocregopher/mediocre-go-lib
1
0
Fork 0
Code Releases Activity

Add miter package

Browse Source
This commit is contained in:
Brian Picciano 2023-12-26 16:39:56 +01:00
parent c76720dffa
commit aac8b11a01
2 changed files with 238 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

136
miter/iter.go Normal file
View File

@ -0,0 +1,136 @@
// Package miter implements a simple iterator type, as well as helper types
// around it.
package miter
import (
"context"
"errors"
)
// ErrEnd is returned from Iterator's Next method when the iterator has
// completed.
var ErrEnd = errors.New("end")
// Iterator is used to visit a sequence of elements one by one.
type Iterator[T any] interface {
// Next will block until:
//
// - It returns the next item in the sequence
// - The Context is canceled and ctx.Err() is returned
// - ErrEnd is returned, indicating there are no more items in the sequence
// - Some other error is returned, indicating an unexpected error occurred
//
// If any non-context error is returned then Next should not be called
// again.
Next(ctx context.Context) (T, error)
}
////////////////////////////////////////////////////////////////////////////////
type iterFn[T any] func(context.Context) (T, error)
// FromFunc wraps a function such that it implements Iterator.
func FromFunc[T any](fn func(context.Context) (T, error)) Iterator[T] {
return iterFn[T](fn)
}
func (f iterFn[T]) Next(ctx context.Context) (T, error) { return f(ctx) }
// Error returns an Iterator which will always return the given error as the
// error result.
func Error[T any](err error) Iterator[T] {
return FromFunc(func(context.Context) (T, error) {
var zero T
return zero, err
})
}
////////////////////////////////////////////////////////////////////////////////
// ToSlice consumes all items off the given Iterator until ErrEnd, creating and
// returning a slice. If the Iterator ever returns any other error then that
// error is returned, along with all consumed items so far.
func ToSlice[T any](ctx context.Context, i Iterator[T]) ([]T, error) {
var s []T
for {
v, err := i.Next(ctx)
if errors.Is(err, ErrEnd) {
return s, nil
} else if err != nil {
return s, err
}
s = append(s, v)
}
}
type iterSlice[T any] []T
// FromSlice wraps a slice such that it implements Iterator.
func FromSlice[T any](s []T) Iterator[T] {
it := iterSlice[T](s)
return &it
}
func (s *iterSlice[T]) Next(context.Context) (T, error) {
if len(*s) == 0 {
var zero T
return zero, ErrEnd
}
next := (*s)[0]
*s = (*s)[1:]
return next, nil
}
////////////////////////////////////////////////////////////////////////////////
type iterCh[T any] <-chan T
// FromChannel wraps a channel such that it implements Iterator.
func FromChannel[T any](ch <-chan T) Iterator[T] {
return iterCh[T](ch)
}
func (ch iterCh[T]) Next(ctx context.Context) (T, error) {
var (
v T
ok bool
)
select {
case v, ok = <-ch:
if !ok {
return v, ErrEnd
}
return v, nil
case <-ctx.Done():
return v, ctx.Err()
}
}
////////////////////////////////////////////////////////////////////////////////
// Lazily returns an Iterator which will initialize itself using the given
// callback upon the first call to Next. If the callback returns an error then
// this error is returned from all calls to Next.
//
// Lazily is useful for cases where an Iterator requires a Context to
// initialize, but no Context is available.
func Lazily[T any](
fn func(ctx context.Context) (Iterator[T], error),
) Iterator[T] {
var i Iterator[T]
return FromFunc(func(ctx context.Context) (T, error) {
if i == nil {
var err error
if i, err = fn(ctx); err != nil {
i = Error[T](err)
}
}
return i.Next(ctx)
})
}

102
miter/util.go Normal file
View File

@ -0,0 +1,102 @@
package miter
import (
"context"
"errors"
)
// Empty returns the empty Iterator, i.e. one which will immediately produce
// ErrEnd.
func Empty[T any]() Iterator[T] { return Error[T](ErrEnd) }
// ForEach calls fn with each item read off the iterator, returning nil once
// ErrEnd is returned from the Iterator or function. If the Iterator or function
// return any other error then that is returned instead.
func ForEach[T any](
ctx context.Context, i Iterator[T], fn func(T) error,
) error {
for {
v, err := i.Next(ctx)
if errors.Is(err, ErrEnd) {
return nil
} else if err != nil {
return err
}
if err := fn(v); errors.Is(err, ErrEnd) {
return nil
} else if err != nil {
return err
}
}
}
// Map will read items from the given Iterator, pass them through the given
// mapping function, and produce the mapped items from the returned Iterator.
//
// If the mapping function returns an error then that is returned from returned
// Iterator, and no more calls should be made to it.
func Map[T1, T2 any](i Iterator[T1], fn func(T1) (T2, error)) Iterator[T2] {
return FromFunc(func(ctx context.Context) (T2, error) {
v, err := i.Next(ctx)
if err != nil {
var zero T2
return zero, err
}
return fn(v)
})
}
// Concat concats all the given Iterators together into a single larger one.
// Each Iterator will be consumed until ErrEnd in turn. Any other errors from
// the inner Iterators will be returned as-is from the outer one.
func Concat[T any](iters ...Iterator[T]) Iterator[T] {
var (
i int
zero T
)
return FromFunc(func(ctx context.Context) (T, error) {
for {
if i >= len(iters) {
return zero, ErrEnd
}
v, err := iters[i].Next(ctx)
if errors.Is(err, ErrEnd) {
i++
continue
}
return v, err
}
})
}
// Filter returns an Iterator which will produce all items from the given
// Iterator for which the function returns true. If the function returns any
// error then that error is returned as-is.
func Filter[T any](
i Iterator[T], fn func(context.Context, T) (bool, error),
) Iterator[T] {
var zero T
return FromFunc(func(ctx context.Context) (T, error) {
for {
v, err := i.Next(ctx)
if err != nil {
return zero, err
}
keep, err := fn(ctx, v)
if err != nil {
return zero, err
}
if keep {
return v, nil
}
}
})
}