Skip to content

Commit

Permalink
Add MustCollect consumer
Browse files Browse the repository at this point in the history
  • Loading branch information
BooleanCat committed Sep 9, 2024
1 parent 0563dae commit e3d53ac
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 3 deletions.
20 changes: 17 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ keys, values := it.Collect2(maps.All(map[string]int{"one": 1, "two": 2}))
keys, values := itx.FromMap(map[string]int{"one": 1, "two": 2}).Collect()
```

### TryCollect
<h3 id="trycollect">TryCollect & MustCollect</h3>

Dealing with iterators that return `T, error` can involve the boilerplate of checking that the
returned slice of errors only contains `nil`. `TryCollect` solves this by collecting all values into
Expand All @@ -90,14 +90,28 @@ if lines, err := it.TryCollect(it.LinesString(text)); err != nil {
}
```

MustCollect is similar except that if an error is encountered then a panic will occur.

```go
text := strings.NewReader("one\ntwo\nthree\n")

lines := it.MustCollect(it.LinesString(text))
```

<!-- prettier-ignore -->
> [!TIP]
> Use `MustCollect` when you can guarantee that no error will occur (such as with
> [strings.Reader](https://pkg.go.dev/strings#Reader)).
<!-- prettier-ignore -->
> [!NOTE]
> If an error is encountered, collection stops. This means the iterator being collected may not be
> fully drained.
<!-- prettier-ignore -->
> [!NOTE]
> The `itx` package does not contain `TryCollect` due to limitations with Go's type system.
> The `itx` package does not contain `TryCollect` or `MustCollect` due to limitations with Go's type
> system.
### ForEach

Expand Down Expand Up @@ -316,7 +330,7 @@ itx.From2(it.Map2(slices.All([]int{1, 2, 3}), printValue2)).Drain()

<!-- prettier-ignore -->
> [!TIP]
> Use Drain to consume an iterator to invoke any side effects when you don't need to collect the
> Use `Drain` to consume an iterator to invoke any side effects when you don't need to collect the
> values.
## Iterators
Expand Down
17 changes: 17 additions & 0 deletions it/iter.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package it

import (
"cmp"
"fmt"
"iter"
)

Expand Down Expand Up @@ -142,6 +143,22 @@ func TryCollect[V any](iterator func(func(V, error) bool)) ([]V, error) {
return values, nil
}

// MustCollect consumes an [iter.Seq2] where the right side yields errors and
// returns a slice of values. If an error is encountered this function will
// panic.
func MustCollect[V any](iterator func(func(V, error) bool)) []V {
var values []V

for v, err := range iterator {
if err != nil {
panic(fmt.Sprintf("it: MustCollect: error yielded by iterator: %s", err.Error()))
}
values = append(values, v)
}

return values
}

// Collect2 consumes an [iter.Seq2] and returns two slices of values.
func Collect2[V, W any](iterator func(func(V, W) bool)) ([]V, []W) {
var (
Expand Down
26 changes: 26 additions & 0 deletions it/iter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -225,3 +225,29 @@ func ExampleDrain2() {
// 2
// 3
}

func ExampleMustCollect() {
buffer := strings.NewReader("one\ntwo")
lines := it.MustCollect(it.LinesString(buffer))

fmt.Println(lines)
// Output: [one two]
}

func TestMustCollectPanic(t *testing.T) {
t.Parallel()

defer func() {
r := recover()

if r == nil {
t.Errorf("expected panic")
}

if fmt.Sprint(r) != "it: MustCollect: error yielded by iterator: read error" {
t.Errorf("wrong panic message")
}
}()

it.MustCollect(it.LinesString(new(failSecondTime)))
}

0 comments on commit e3d53ac

Please sign in to comment.