Add take()

This commit is contained in:
Anton Medvedev
2023-08-30 12:54:13 +02:00
parent 7d18813005
commit 66706467ba
2 changed files with 39 additions and 0 deletions
+36
View File
@@ -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) {
+3
View File
@@ -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},
@@ -123,6 +125,7 @@ func TestBuiltin_works_with_any(t *testing.T) {
arity int
}{
"get": {2},
"take": {2},
"sortBy": {2},
}