pkg/list: adding flattenN function as a builtin
this commit adds a new builtin to the list package for generating a
flat list from a nested list by expanding nested lists in place. only
nested lists n levels deep into the structure are expanded.
for instance
list.FlattenN([1, [[2, 3], []], [4]], 1)
results in
[1, [2, 3], [], 4]
Change-Id: I991caf51b90907783e25c780d50857c605b46a8d
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/3720
Reviewed-by: Felix Ehrenpfort <felix.ehrenpfort@codecentric.de>
Reviewed-by: Marcel van Lohuizen <mpvl@golang.org>
diff --git a/cue/builtin_test.go b/cue/builtin_test.go
index 56c044f..7e5d41e 100644
--- a/cue/builtin_test.go
+++ b/cue/builtin_test.go
@@ -162,6 +162,24 @@
test("list", `list.Flatten("foo")`),
`_|_(error in call to list.Flatten: cannot use value "foo" (type string) as list)`,
}, {
+ test("list", `list.FlattenN([1, [[2, 3], []], [4]], -1)`),
+ `[1,2,3,4]`,
+ }, {
+ test("list", `list.FlattenN([1, [[2, 3], []], [4]], 0)`),
+ `[1,[[2,3],[]],[4]]`,
+ }, {
+ test("list", `list.FlattenN([1, [[2, 3], []], [4]], 1)`),
+ `[1,[2,3],[],4]`,
+ }, {
+ test("list", `list.FlattenN([1, [[2, 3], []], [4]], 2)`),
+ `[1,2,3,4]`,
+ }, {
+ test("list", `list.FlattenN("foo", 1)`),
+ `_|_(error in call to list.FlattenN: cannot use value "foo" (type string) as list)`,
+ }, {
+ test("list", `list.FlattenN([], "foo")`),
+ `_|_(cannot use "foo" (type string) as number in argument 2 to list.FlattenN)`,
+ }, {
test("list", `list.Max([1, 2, 3, 4])`),
`4`,
}, {
diff --git a/cue/builtins.go b/cue/builtins.go
index e5f9044..ef2f823 100644
--- a/cue/builtins.go
+++ b/cue/builtins.go
@@ -717,6 +717,43 @@
}()
},
}, {
+ Name: "FlattenN",
+ Params: []kind{topKind, numKind},
+ Result: listKind,
+ Func: func(c *callCtxt) {
+ xs, depth := c.value(0), c.decimal(1)
+ c.ret, c.err = func() (interface{}, error) {
+ var flattenN func(Value, *internal.Decimal) ([]Value, error)
+ one := apd.New(1, 0)
+ flattenN = func(xs Value, depth *internal.Decimal) ([]Value, error) {
+ var res []Value
+ iter, err := xs.List()
+ if err != nil {
+ return nil, err
+ }
+ for iter.Next() {
+ val := iter.Value()
+ if val.Kind() == ListKind && !depth.IsZero() {
+ d := apd.New(0, 0)
+ _, err := internal.BaseContext.Sub(d, depth, one)
+ if err != nil {
+ return nil, err
+ }
+ vals, err := flattenN(val, d)
+ if err != nil {
+ return nil, err
+ }
+ res = append(res, vals...)
+ } else {
+ res = append(res, val)
+ }
+ }
+ return res, nil
+ }
+ return flattenN(xs, depth)
+ }()
+ },
+ }, {
Name: "Take",
Params: []kind{listKind, intKind},
Result: listKind,
diff --git a/pkg/list/list.go b/pkg/list/list.go
index e6b0dff..7f37933 100644
--- a/pkg/list/list.go
+++ b/pkg/list/list.go
@@ -20,6 +20,8 @@
"sort"
"cuelang.org/go/cue"
+ "cuelang.org/go/internal"
+ "github.com/cockroachdb/apd/v2"
)
// Drop reports the suffix of list x after the first n elements,
@@ -45,7 +47,7 @@
return x[n:], nil
}
-// Flatten reports a flattend sequence of the list x by expanding any elements
+// Flatten reports a flattend sequence of the list xs by expanding any elements
// that are lists.
//
// For instance:
@@ -81,6 +83,48 @@
return flatten(xs)
}
+// FlattenN reports a flattend sequence of the list xs by expanding any elements
+// depth levels deep. If depth is negative all elements are expanded.
+//
+// For instance:
+//
+// FlattenN([1, [[2, 3], []], [4]], 1)
+//
+// results in
+//
+// [1, [2, 3], [], 4]
+//
+func FlattenN(xs cue.Value, depth *internal.Decimal) ([]cue.Value, error) {
+ var flattenN func(cue.Value, *internal.Decimal) ([]cue.Value, error)
+ one := apd.New(1, 0)
+ flattenN = func(xs cue.Value, depth *internal.Decimal) ([]cue.Value, error) {
+ var res []cue.Value
+ iter, err := xs.List()
+ if err != nil {
+ return nil, err
+ }
+ for iter.Next() {
+ val := iter.Value()
+ if val.Kind() == cue.ListKind && !depth.IsZero() {
+ d := apd.New(0, 0)
+ _, err := internal.BaseContext.Sub(d, depth, one)
+ if err != nil {
+ return nil, err
+ }
+ vals, err := flattenN(val, d)
+ if err != nil {
+ return nil, err
+ }
+ res = append(res, vals...)
+ } else {
+ res = append(res, val)
+ }
+ }
+ return res, nil
+ }
+ return flattenN(xs, depth)
+}
+
// Take reports the prefix of length n of list x, or x itself if n > len(x).
//
// For instance:
