Add take()

Author: antonmedv

builtin/builtin.go

@@ -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) {

builtin/builtin_test.go

@@ -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},
 	}