feat: add SampleBy and SamplesBy (#516)

* Add SampleBy and SamplesBy

Co-authored-by: Bram Van de Walle <[email protected]>
Co-authored-by: Samuel Berthe <[email protected]>

README.md

@@ -249,7 +249,9 @@ Supported search helpers:
 - [LastOr](#LastOr)
 - [Nth](#nth)
 - [Sample](#sample)
+- [SampleBy](#sampleby)
 - [Samples](#samples)
+- [SamplesBy](#samplesby)
 
 Conditional helpers:
 
@@ -2716,6 +2718,21 @@ lo.Sample([]string{})
 // ""
 ```
 
+### SampleBy
+
+Returns a random item from collection, using a given random integer generator.
+
+```go
+import "math/rand"
+
+r := rand.New(rand.NewSource(42))
+lo.SampleBy([]string{"a", "b", "c"}, r.Intn)
+// a random string from []string{"a", "b", "c"}, using a seeded random generator
+
+lo.SampleBy([]string{}, r.Intn)
+// ""
+```
+
 ### Samples
 
 Returns N random unique items from collection.
@@ -2725,6 +2742,16 @@ lo.Samples([]string{"a", "b", "c"}, 3)
 // []string{"a", "b", "c"} in random order
 ```
 
+### SamplesBy
+
+Returns N random unique items from collection, using a given random integer generator.
+
+```go
+r := rand.New(rand.NewSource(42))
+lo.SamplesBy([]string{"a", "b", "c"}, 3, r.Intn)
+// []string{"a", "b", "c"} in random order, using a seeded random generator
+```
+
 ### Ternary
 
 A 1 line if/else statement.

find.go

@@ -579,18 +579,33 @@ func Nth[T any, N constraints.Integer](collection []T, nth N) (T, error) {
 	return collection[l+n], nil
 }
 
+// randomIntGenerator is a function that should return a random integer in the range [0, n)
+// where n is the parameter passed to the randomIntGenerator.
+type randomIntGenerator func(n int) int
+
 // Sample returns a random item from collection.
 func Sample[T any](collection []T) T {
+	result := SampleBy(collection, rand.IntN)
+	return result
+}
+
+// SampleBy returns a random item from collection, using randomIntGenerator as the random index generator.
+func SampleBy[T any](collection []T, randomIntGenerator randomIntGenerator) T {
 	size := len(collection)
 	if size == 0 {
 		return Empty[T]()
 	}
-
-	return collection[rand.IntN(size)]
+	return collection[randomIntGenerator(size)]
 }
 
 // Samples returns N random unique items from collection.
 func Samples[T any, Slice ~[]T](collection Slice, count int) Slice {
+	results := SamplesBy(collection, count, rand.IntN)
+	return results
+}
+
+// SamplesBy returns N random unique items from collection, using randomIntGenerator as the random index generator.
+func SamplesBy[T any, Slice ~[]T](collection Slice, count int, randomIntGenerator randomIntGenerator) Slice {
 	size := len(collection)
 
 	copy := append(Slice{}, collection...)
@@ -600,7 +615,7 @@ func Samples[T any, Slice ~[]T](collection Slice, count int) Slice {
 	for i := 0; i < size && i < count; i++ {
 		copyLength := size - i
 
-		index := rand.IntN(size - i)
+		index := randomIntGenerator(size - i)
 		results = append(results, copy[index])
 
 		// Removes element.

find_test.go

@@ -643,6 +643,19 @@ func TestSample(t *testing.T) {
 	is.Equal(result2, "")
 }
 
+func TestSampleBy(t *testing.T) {
+	t.Parallel()
+	is := assert.New(t)
+
+	r := rand.New(rand.NewSource(42))
+
+	result1 := SampleBy([]string{"a", "b", "c"}, r.Intn)
+	result2 := SampleBy([]string{}, rand.Intn)
+
+	is.True(Contains([]string{"a", "b", "c"}, result1))
+	is.Equal(result2, "")
+}
+
 func TestSamples(t *testing.T) {
 	t.Parallel()
 	is := assert.New(t)
@@ -662,3 +675,23 @@ func TestSamples(t *testing.T) {
 	nonempty := Samples(allStrings, 2)
 	is.IsType(nonempty, allStrings, "type preserved")
 }
+
+func TestSamplesBy(t *testing.T) {
+	t.Parallel()
+	is := assert.New(t)
+
+	r := rand.New(rand.NewSource(42))
+
+	result1 := SamplesBy([]string{"a", "b", "c"}, 3, r.Intn)
+	result2 := SamplesBy([]string{}, 3, r.Intn)
+
+	sort.Strings(result1)
+
+	is.Equal(result1, []string{"a", "b", "c"})
+	is.Equal(result2, []string{})
+
+	type myStrings []string
+	allStrings := myStrings{"", "foo", "bar"}
+	nonempty := SamplesBy(allStrings, 2, r.Intn)
+	is.IsType(nonempty, allStrings, "type preserved")
+}