Skip to content

Commit

Permalink
feat: Implement funk.Partition
Browse files Browse the repository at this point in the history
  • Loading branch information
nguyenvantuan2391996 committed Aug 8, 2023
1 parent 045ef11 commit 5cdac4c
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 0 deletions.
48 changes: 48 additions & 0 deletions partition.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
package funk

import (
"reflect"
)

// Partition separates the elements of the input array into two slices based on the predicate function.
// It takes an array-like data structure and a predicate function that determines the Partition.
// The predicate function should have the signature func(elementType) bool.
// The function returns two new slices: one containing elements that satisfy the predicate,
// and the other containing elements that do not satisfy the predicate.
func Partition(in, predicate interface{}) interface{} {
if !IsCollection(in) {
panic("First parameter must be a collection")
}

if !IsFunction(predicate, 1, 1) {
panic("Second argument must be function")
}

inValue, funcValue := reflect.ValueOf(in), reflect.ValueOf(predicate)

funcType := funcValue.Type()

if funcType.Out(0).Kind() != reflect.Bool {
panic("Return argument should be a boolean")
}

result := reflect.MakeSlice(reflect.SliceOf(reflect.TypeOf(in)), 0, 0)
partitionOne, partitionTwo := reflect.MakeSlice(inValue.Type(), 0, 0), reflect.MakeSlice(inValue.Type(), 0, 0)

for i := 0; i < inValue.Len(); i++ {
element := inValue.Index(i)

res := funcValue.Call([]reflect.Value{reflect.ValueOf(element.Interface())})
if res[0].Interface().(bool) {
partitionOne = reflect.Append(partitionOne, element)
} else {
partitionTwo = reflect.Append(partitionTwo, element)
}
}

if partitionOne.Len() > 0 || partitionTwo.Len() > 0 {
result = reflect.Append(result, partitionOne, partitionTwo)
}

return result.Interface()
}
35 changes: 35 additions & 0 deletions partition_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
package funk

import (
"testing"

"github.com/stretchr/testify/assert"
)

func Test_partition_valid_int(t *testing.T) {
result := Partition([]int{1, 2, 3, 4}, func(n int) bool {
return n%2 == 0
})

assert.Equal(t, [][]int{{2, 4}, {1, 3}}, result)
}

func Test_partition_valid_string(t *testing.T) {
result := Partition([]string{"a", "b", "c"}, func(n string) bool {
return n > "a"
})

assert.Equal(t, [][]string{{"b", "c"}, {"a"}}, result)
}

func Test_partition_valid_struct(t *testing.T) {
type User struct {
Name string
Age int
}
result := Partition([]*User{{Name: "Kakalot", Age: 26}, {Name: "Vegeta", Age: 27}, {Name: "Trunk", Age: 10}}, func(n *User) bool {
return n.Age%2 == 0
})

assert.Equal(t, [][]*User{{{Name: "Kakalot", Age: 26}, {Name: "Trunk", Age: 10}}, {{Name: "Vegeta", Age: 27}}}, result)
}

0 comments on commit 5cdac4c

Please sign in to comment.