Skip to content

Commit

Permalink
Add take()
Browse files Browse the repository at this point in the history
  • Loading branch information
antonmedv committed Aug 30, 2023
1 parent 418fed5 commit 82bb7df
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 0 deletions.
36 changes: 36 additions & 0 deletions builtin/builtin.go
Original file line number Diff line number Diff line change
Expand Up @@ -524,6 +524,42 @@ var Builtins = []*ast.Function{
return runtime.Fetch(args[0], args[1]), nil
},
},
{
Name: "take",
Func: func(args ...any) (any, error) {
if len(args) != 2 {
return nil, fmt.Errorf("invalid number of arguments (expected 2, got %d)", len(args))
}
v := reflect.ValueOf(args[0])
if v.Kind() != reflect.Slice && v.Kind() != reflect.Array {
return nil, fmt.Errorf("cannot take from %s", v.Kind())
}
n := reflect.ValueOf(args[1])
if !n.CanInt() {
return nil, fmt.Errorf("cannot take %s elements", n.Kind())
}
if n.Int() > int64(v.Len()) {
return args[0], nil
}
return v.Slice(0, int(n.Int())).Interface(), nil
},
Validate: func(args []reflect.Type) (reflect.Type, error) {
if len(args) != 2 {
return anyType, fmt.Errorf("invalid number of arguments (expected 2, got %d)", len(args))
}
switch kind(args[0]) {
case reflect.Interface, reflect.Slice, reflect.Array:
default:
return anyType, fmt.Errorf("cannot take from %s", args[0])
}
switch kind(args[1]) {
case reflect.Interface, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
default:
return anyType, fmt.Errorf("cannot take %s elements", args[1])
}
return args[0], nil
},
},
{
Name: "keys",
Func: func(args ...any) (any, error) {
Expand Down
3 changes: 3 additions & 0 deletions builtin/builtin_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,8 @@ func TestBuiltin(t *testing.T) {
{`get(ArrayOfAny, 1)`, "2"},
{`get({foo: 1, bar: 2}, "foo")`, 1},
{`get({foo: 1, bar: 2}, "unknown")`, nil},
{`take(ArrayOfString, 2)`, []string{"foo", "bar"}},
{`take(ArrayOfString, 99)`, []string{"foo", "bar", "baz"}},
{`"foo" in keys({foo: 1, bar: 2})`, true},
{`1 in values({foo: 1, bar: 2})`, true},
{`len(toPairs({foo: 1, bar: 2}))`, 2},
Expand Down Expand Up @@ -123,6 +125,7 @@ func TestBuiltin_works_with_any(t *testing.T) {
arity int
}{
"get": {2},
"take": {2},
"sortBy": {2},
}

Expand Down

0 comments on commit 82bb7df

Please # to comment.