diff --git a/moslice/batch.go b/moslice/batch.go new file mode 100644 index 0000000..79d2350 --- /dev/null +++ b/moslice/batch.go @@ -0,0 +1,20 @@ +package moslice + +func Batch[T any](ts []T, size int) [][]T { + if len(ts) == 0 || size <= 0 { + return nil + } + + numBatches := (len(ts) + size - 1) / size + batches := make([][]T, 0, numBatches) + + for i := 0; i < len(ts); i += size { + end := i + size + if end > len(ts) { + end = len(ts) + } + batches = append(batches, ts[i:end]) + } + + return batches +} diff --git a/moslice/batch_test.go b/moslice/batch_test.go new file mode 100644 index 0000000..f80c17b --- /dev/null +++ b/moslice/batch_test.go @@ -0,0 +1,52 @@ +package moslice + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestBatch(t *testing.T) { + t.Run("empty slice", func(t *testing.T) { + result := Batch([]int{}, 3) + assert.Nil(t, result) + }) + + t.Run("zero or negative size", func(t *testing.T) { + result := Batch([]int{1, 2, 3}, 0) + assert.Nil(t, result) + + result = Batch([]int{1, 2, 3}, -1) + assert.Nil(t, result) + }) + + t.Run("slice smaller than batch size", func(t *testing.T) { + result := Batch([]int{1, 2}, 5) + assert.Equal(t, [][]int{{1, 2}}, result) + }) + + t.Run("slice equal to batch size", func(t *testing.T) { + result := Batch([]int{1, 2, 3}, 3) + assert.Equal(t, [][]int{{1, 2, 3}}, result) + }) + + t.Run("slice evenly divisible by batch size", func(t *testing.T) { + result := Batch([]int{1, 2, 3, 4, 5, 6}, 2) + assert.Equal(t, [][]int{{1, 2}, {3, 4}, {5, 6}}, result) + }) + + t.Run("slice not evenly divisible by batch size", func(t *testing.T) { + result := Batch([]int{1, 2, 3, 4, 5}, 2) + assert.Equal(t, [][]int{{1, 2}, {3, 4}, {5}}, result) + }) + + t.Run("batch size of 1", func(t *testing.T) { + result := Batch([]int{1, 2, 3}, 1) + assert.Equal(t, [][]int{{1}, {2}, {3}}, result) + }) + + t.Run("string slice", func(t *testing.T) { + result := Batch([]string{"a", "b", "c", "d", "e"}, 2) + assert.Equal(t, [][]string{{"a", "b"}, {"c", "d"}, {"e"}}, result) + }) +}