diff --git a/README.md b/README.md index eddbdecb..5019b890 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,8 @@ Supported helpers for slices: - [Compact](#compact) - [IsSorted](#issorted) - [IsSortedByKey](#issortedbykey) +- [Insert](#insert) +- [InsertSlice](#insertslice) Supported helpers for maps: @@ -923,6 +925,28 @@ slice := lo.IsSortedByKey([]string{"a", "bb", "ccc"}, func(s string) int { [[play](https://go.dev/play/p/wiG6XyBBu49)] +### Insert + +Inserts an element at index i. + +```go +result := lo.Insert(["a", "c", "d"], 1, "b") +// []string{"a", "b", "c", "d"} +``` + +[[play](https://go.dev/play/p/Eu3pk6_ODVQ)] + +### InsertSlice + +Inserts multiple elements at index i. + +```go +result := lo.InsertSlice([]string{"a", "d"}, 1, []string{"b", "c"}) +// []string{"a", "b", "c", "d"} +``` + +[[play](https://go.dev/play/p/8gXdkTovsIC)] + ### Keys Creates an array of the map keys. diff --git a/slice.go b/slice.go index 49c991f8..5986aa03 100644 --- a/slice.go +++ b/slice.go @@ -592,3 +592,22 @@ func IsSortedByKey[T any, K constraints.Ordered](collection []T, iteratee func(i return true } + +// Insert inserts an element at index i. +// Play: https://go.dev/play/p/Eu3pk6_ODVQ +func Insert[T any](collection []T, i int, element T) []T { + return append(collection[:i], append([]T{element}, collection[i:]...)...) +} + +// InsertSlice inserts multiple elements at index i. +// Play: https://go.dev/play/p/8gXdkTovsIC +func InsertSlice[T any](collection []T, i int, elements []T) []T { + if i < 0 || i > len(collection) { + panic("index out of bounds") + } + if len(elements) == 0 { + return collection + } + + return append(append(collection[:i], elements...), collection[i:]...) +} diff --git a/slice_test.go b/slice_test.go index 03326a01..f9869cd8 100644 --- a/slice_test.go +++ b/slice_test.go @@ -388,7 +388,7 @@ func TestAssociate(t *testing.T) { func TestSliceToMap(t *testing.T) { t.Parallel() - + type foo struct { baz string bar int @@ -626,7 +626,7 @@ func TestSlice(t *testing.T) { out16 := Slice(in, -10, 1) out17 := Slice(in, -1, 3) out18 := Slice(in, -10, 7) - + is.Equal([]int{}, out1) is.Equal([]int{0}, out2) is.Equal([]int{0, 1, 2, 3, 4}, out3) @@ -759,3 +759,25 @@ func TestIsSortedByKey(t *testing.T) { return ret })) } + +func TestInsert(t *testing.T) { + t.Parallel() + is := assert.New(t) + + inserts := Insert([]string{"a", "c", "d"}, 1, "b") + is.True(len(inserts) == 4 && inserts[1] == "b") + + inserti := Insert([]int{1, 3, 4}, 1, 2) + is.True(len(inserti) == 4 && inserti[1] == 2) +} + +func TestInsertSlice(t *testing.T) { + t.Parallel() + is := assert.New(t) + + inserts := InsertSlice([]string{"a", "d"}, 1, []string{"b", "c"}) + is.True(len(inserts) == 4 && inserts[1] == "b" && inserts[2] == "c") + + inserti := InsertSlice([]int{1, 4}, 1, []int{2, 3}) + is.True(len(inserti) == 4 && inserti[1] == 2 && inserti[2] == 3) +}