Skip to content

Commit e3d53ac

Browse files
committed
Add MustCollect consumer
1 parent 0563dae commit e3d53ac

File tree

3 files changed

+60
-3
lines changed

3 files changed

+60
-3
lines changed

README.md

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ keys, values := it.Collect2(maps.All(map[string]int{"one": 1, "two": 2}))
7676
keys, values := itx.FromMap(map[string]int{"one": 1, "two": 2}).Collect()
7777
```
7878

79-
### TryCollect
79+
<h3 id="trycollect">TryCollect & MustCollect</h3>
8080

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

93+
MustCollect is similar except that if an error is encountered then a panic will occur.
94+
95+
```go
96+
text := strings.NewReader("one\ntwo\nthree\n")
97+
98+
lines := it.MustCollect(it.LinesString(text))
99+
```
100+
101+
<!-- prettier-ignore -->
102+
> [!TIP]
103+
> Use `MustCollect` when you can guarantee that no error will occur (such as with
104+
> [strings.Reader](https://pkg.go.dev/strings#Reader)).
105+
93106
<!-- prettier-ignore -->
94107
> [!NOTE]
95108
> If an error is encountered, collection stops. This means the iterator being collected may not be
96109
> fully drained.
97110
98111
<!-- prettier-ignore -->
99112
> [!NOTE]
100-
> The `itx` package does not contain `TryCollect` due to limitations with Go's type system.
113+
> The `itx` package does not contain `TryCollect` or `MustCollect` due to limitations with Go's type
114+
> system.
101115
102116
### ForEach
103117

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

317331
<!-- prettier-ignore -->
318332
> [!TIP]
319-
> Use Drain to consume an iterator to invoke any side effects when you don't need to collect the
333+
> Use `Drain` to consume an iterator to invoke any side effects when you don't need to collect the
320334
> values.
321335
322336
## Iterators

it/iter.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package it
22

33
import (
44
"cmp"
5+
"fmt"
56
"iter"
67
)
78

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

146+
// MustCollect consumes an [iter.Seq2] where the right side yields errors and
147+
// returns a slice of values. If an error is encountered this function will
148+
// panic.
149+
func MustCollect[V any](iterator func(func(V, error) bool)) []V {
150+
var values []V
151+
152+
for v, err := range iterator {
153+
if err != nil {
154+
panic(fmt.Sprintf("it: MustCollect: error yielded by iterator: %s", err.Error()))
155+
}
156+
values = append(values, v)
157+
}
158+
159+
return values
160+
}
161+
145162
// Collect2 consumes an [iter.Seq2] and returns two slices of values.
146163
func Collect2[V, W any](iterator func(func(V, W) bool)) ([]V, []W) {
147164
var (

it/iter_test.go

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,3 +225,29 @@ func ExampleDrain2() {
225225
// 2
226226
// 3
227227
}
228+
229+
func ExampleMustCollect() {
230+
buffer := strings.NewReader("one\ntwo")
231+
lines := it.MustCollect(it.LinesString(buffer))
232+
233+
fmt.Println(lines)
234+
// Output: [one two]
235+
}
236+
237+
func TestMustCollectPanic(t *testing.T) {
238+
t.Parallel()
239+
240+
defer func() {
241+
r := recover()
242+
243+
if r == nil {
244+
t.Errorf("expected panic")
245+
}
246+
247+
if fmt.Sprint(r) != "it: MustCollect: error yielded by iterator: read error" {
248+
t.Errorf("wrong panic message")
249+
}
250+
}()
251+
252+
it.MustCollect(it.LinesString(new(failSecondTime)))
253+
}

0 commit comments

Comments
 (0)