Merge cc7940d into 417b4e8

docs/types/array.md

@@ -423,6 +423,15 @@ Flattens an array a single level deep:
 [[1, 2, 3, 4]].flatten() # [1, 2, 3, 4]
 ```
 
+### flatten_deep()
+
+Flattens an array recursively until no member is an array:
+
+```py
+[[[1, 2], [[[[3]]]], [4]]].flatten_deep() # [1, 2, 3, 4]
+[[1, [2, 3], 4]].flatten_deep() # [1, 2, 3, 4]
+```
+
 ## Next
 
 That's about it for this section!

evaluator/builtin_functions_test.go

@@ -678,6 +678,19 @@ func TestFlatten(t *testing.T) {
 	testBuiltinFunction(tests, t)
 }
 
+func TestFlattenDeep(t *testing.T) {
+	tests := []Tests{
+		{`[1, 2, 3].flatten_deep()`, []int{1, 2, 3}},
+		{`[1, 2, [3]].flatten_deep()`, []int{1, 2, 3}},
+		{`[1, 2, [3, 4]].flatten_deep()`, []int{1, 2, 3, 4}},
+		{`[[1, 2], [3, 4]].flatten_deep()`, []int{1, 2, 3, 4}},
+		{`[[1, [2]], [3, 4]].flatten_deep()`, []int{1, 2, 3, 4}},
+		{`[[[1, [2]], [3, 4]]].flatten_deep()`, []int{1, 2, 3, 4}},
+	}
+
+	testBuiltinFunction(tests, t)
+}
+
 func testBuiltinFunction(tests []Tests, t *testing.T) {
 	for _, tt := range tests {
 		evaluated := testEval(tt.input)

evaluator/functions.go

@@ -223,6 +223,11 @@ func getFns() map[string]*object.Builtin {
 			Types: []string{object.ARRAY_OBJ},
 			Fn:    flattenFn,
 		},
+		// flattem(array:[1, 2, 3])
+		"flatten_deep": &object.Builtin{
+			Types: []string{object.ARRAY_OBJ},
+			Fn:    flattenDeepFn,
+		},
 		// map(array:[1, 2, 3], function:f(x) { x + 1 })
 		"map": &object.Builtin{
 			Types: []string{object.ARRAY_OBJ},
@@ -1297,6 +1302,15 @@ func unionFn(tok token.Token, env *object.Environment, args ...object.Object) ob
 
 // flatten(array:[1, 2, 3])
 func flattenFn(tok token.Token, env *object.Environment, args ...object.Object) object.Object {
+	return flatten("flatten", false, tok, env, args...)
+}
+
+// flatten_deep(array:[1, 2, 3])
+func flattenDeepFn(tok token.Token, env *object.Environment, args ...object.Object) object.Object {
+	return flatten("flatten_deep", true, tok, env, args...)
+}
+
+func flatten(fnName string, deep bool, tok token.Token, env *object.Environment, args ...object.Object) object.Object {
 	err := validateArgs(tok, "flatten", args, 1, [][]string{{object.ARRAY_OBJ}})
 	if err != nil {
 		return err
@@ -1308,8 +1322,12 @@ func flattenFn(tok token.Token, env *object.Environment, args ...object.Object)
 	for _, v := range originalElements {
 		switch e := v.(type) {
 		case *object.Array:
-			for _, x := range e.Elements {
-				elements = append(elements, x)
+			if deep {
+				elements = append(elements, flattenDeepFn(tok, env, e).(*object.Array).Elements...)
+			} else {
+				for _, x := range e.Elements {
+					elements = append(elements, x)
+				}
 			}
 		default:
 			elements = append(elements, e)