Merge branch 'nicklaus-dev-master'

README.md

@@ -102,6 +102,7 @@ Supported helpers for slices:
 - [RepeatBy](#repeatby)
 - [KeyBy](#keyby)
 - [Associate / SliceToMap](#associate-alias-slicetomap)
+- [Keyify](#keyify)
 - [Drop](#drop)
 - [DropRight](#dropright)
 - [DropWhile](#dropwhile)
@@ -191,6 +192,7 @@ Supported helpers for channels:
 - [SliceToChannel](#slicetochannel)
 - [Generator](#generator)
 - [Buffer](#buffer)
+- [BufferWithContext](#bufferwithcontext)
 - [BufferWithTimeout](#bufferwithtimeout)
 - [FanIn](#fanin)
 - [FanOut](#fanout)
@@ -762,6 +764,15 @@ aMap := lo.Associate(in, func (f *foo) (string, int) {
 
 [[play](https://go.dev/play/p/WHa2CfMO3Lr)]
 
+### Keyify
+
+Returns a map with each unique element of the slice as a key.
+
+```go
+set := lo.Keyify([]int{1, 1, 2, 3, 4})
+// map[int]struct{}{1:{}, 2:{}, 3:{}, 4:{}}
+```
+
 ### Drop
 
 Drops n elements from the beginning of a slice or array.
@@ -1928,6 +1939,32 @@ for {
 }
 ```
 
+### BufferWithContext
+
+Creates a slice of n elements from a channel, with timeout. Returns the slice, the slice length, the read time and the channel status (opened/closed).
+
+```go
+ctx, cancel := context.WithCancel(context.TODO())
+go func() {
+    ch <- 0
+    time.Sleep(10*time.Millisecond)
+    ch <- 1
+    time.Sleep(10*time.Millisecond)
+    ch <- 2
+    time.Sleep(10*time.Millisecond)
+    ch <- 3
+    time.Sleep(10*time.Millisecond)
+    ch <- 4
+    time.Sleep(10*time.Millisecond)
+    cancel()
+}()
+
+items1, length1, duration1, ok1 := lo.BufferWithContext(ctx, ch, 3)
+// []int{0, 1, 2}, 3, 20ms, true
+items2, length2, duration2, ok2 := lo.BufferWithContext(ctx, ch, 3)
+// []int{3, 4}, 2, 30ms, false
+```
+
 ### BufferWithTimeout
 
 Creates a slice of n elements from a channel, with timeout. Returns the slice, the slice length, the read time and the channel status (opened/closed).

channel.go

@@ -1,6 +1,7 @@
 package lo
 
 import (
+	"context"
 	"sync"
 	"time"
 
@@ -220,17 +221,13 @@ func Batch[T any](ch <-chan T, size int) (collection []T, length int, readTime t
 	return Buffer(ch, size)
 }
 
-// BufferWithTimeout creates a slice of n elements from a channel, with timeout. Returns the slice and the slice length.
+// BufferWithContext creates a slice of n elements from a channel, with context. Returns the slice and the slice length.
 // @TODO: we should probably provide an helper that reuse the same buffer.
-func BufferWithTimeout[T any](ch <-chan T, size int, timeout time.Duration) (collection []T, length int, readTime time.Duration, ok bool) {
-	expire := time.NewTimer(timeout)
-	defer expire.Stop()
-
+func BufferWithContext[T any](ctx context.Context, ch <-chan T, size int) (collection []T, length int, readTime time.Duration, ok bool) {
 	buffer := make([]T, 0, size)
-	index := 0
 	now := time.Now()
 
-	for ; index < size; index++ {
+	for index := 0; index < size; index++ {
 		select {
 		case item, ok := <-ch:
 			if !ok {
@@ -239,12 +236,19 @@ func BufferWithTimeout[T any](ch <-chan T, size int, timeout time.Duration) (col
 
 			buffer = append(buffer, item)
 
-		case <-expire.C:
+		case <-ctx.Done():
 			return buffer, index, time.Since(now), true
 		}
 	}
 
-	return buffer, index, time.Since(now), true
+	return buffer, size, time.Since(now), true
+}
+
+// BufferWithTimeout creates a slice of n elements from a channel, with timeout. Returns the slice and the slice length.
+func BufferWithTimeout[T any](ch <-chan T, size int, timeout time.Duration) (collection []T, length int, readTime time.Duration, ok bool) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	return BufferWithContext(ctx, ch, size)
 }
 
 // BatchWithTimeout creates a slice of n elements from a channel, with timeout. Returns the slice and the slice length.

channel_test.go

@@ -1,6 +1,7 @@
 package lo
 
 import (
+	"context"
 	"math/rand"
 	"testing"
 	"time"
@@ -279,6 +280,42 @@ func TestBuffer(t *testing.T) {
 	is.False(ok3)
 }
 
+func TestBufferWithContext(t *testing.T) {
+	t.Parallel()
+	testWithTimeout(t, 200*time.Millisecond)
+	is := assert.New(t)
+
+	ch1 := make(chan int, 10)
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		ch1 <- 0
+		ch1 <- 1
+		ch1 <- 2
+		time.Sleep(5 * time.Millisecond)
+		cancel()
+		ch1 <- 3
+		ch1 <- 4
+		ch1 <- 5
+		close(ch1)
+	}()
+	items1, length1, _, ok1 := BufferWithContext(ctx, ch1, 20)
+	is.Equal([]int{0, 1, 2}, items1)
+	is.Equal(3, length1)
+	is.True(ok1)
+
+	ch2 := make(chan int, 10)
+	ctx, cancel = context.WithCancel(context.Background())
+	defer cancel()
+	defer close(ch2)
+	for i := 0; i < 10; i++ {
+		ch2 <- i
+	}
+	items2, length2, _, ok2 := BufferWithContext(ctx, ch2, 5)
+	is.Equal([]int{0, 1, 2, 3, 4}, items2)
+	is.Equal(5, length2)
+	is.True(ok2)
+}
+
 func TestBufferWithTimeout(t *testing.T) {
 	t.Parallel()
 	testWithTimeout(t, 200*time.Millisecond)

slice.go

@@ -388,6 +388,17 @@ func SliceToMap[T any, K comparable, V any](collection []T, transform func(item
 	return Associate(collection, transform)
 }
 
+// Keyify returns a map with each unique element of the slice as a key.
+func Keyify[T comparable, Slice ~[]T](collection Slice) map[T]struct{} {
+	result := make(map[T]struct{}, len(collection))
+
+	for _, item := range collection {
+		result[item] = struct{}{}
+	}
+
+	return result
+}
+
 // Drop drops n elements from the beginning of a slice or array.
 // Play: https://go.dev/play/p/JswS7vXRJP2
 func Drop[T any, Slice ~[]T](collection Slice, n int) Slice {

slice_example_test.go

@@ -287,6 +287,22 @@ func ExampleAssociate() {
 	// Output: map[a:1 aa:2 aaa:3]
 }
 
+func ExampleKeyify() {
+	list := []string{"a", "a", "b", "b", "d"}
+
+	set := Keyify(list)
+	_, ok1 := set["a"]
+	_, ok2 := set["c"]
+	fmt.Printf("%v\n", ok1)
+	fmt.Printf("%v\n", ok2)
+	fmt.Printf("%v\n", set)
+
+	// Output:
+	// true
+	// false
+	// map[a:{} b:{} d:{}]
+}
+
 func ExampleDrop() {
 	list := []int{0, 1, 2, 3, 4, 5}
 

slice_test.go

@@ -531,6 +531,18 @@ func TestSliceToMap(t *testing.T) {
 	}
 }
 
+func TestKeyify(t *testing.T) {
+	t.Parallel()
+	is := assert.New(t)
+
+	result1 := Keyify([]int{1, 2, 3, 4})
+	result2 := Keyify([]int{1, 1, 1, 2})
+	result3 := Keyify([]int{})
+	is.Equal(result1, map[int]struct{}{1: {}, 2: {}, 3: {}, 4: {}})
+	is.Equal(result2, map[int]struct{}{1: {}, 2: {}})
+	is.Equal(result3, map[int]struct{}{})
+}
+
 func TestDrop(t *testing.T) {
 	t.Parallel()
 	is := assert.New(t)