fix: Fix error messages when a comprehension iterates over a non-array.

rudo-thomas · johnbartholomew · commit 1c79f577ba95 · 2025-03-21T16:37:59.000Z
Problems:
 - When iterating over an empty string in a list comprehension, the
   result is an empty string. This is a bug, it should be an error.
 - When iterating over a non-empty string in a list comprehension, the
   expected and unexpected types in the error message are swapped.
 - Error messages mention "std.flatMap" when object/list comprehensions
   would iterate over a value that is neither array nor string.

```
$ jsonnet --version
Jsonnet commandline interpreter (Go implementation) v0.21.0-rc2
$ jsonnet -e '[a for a in ""]'
""
$ jsonnet -e '[a for a in "b"]'
RUNTIME ERROR: Unexpected type array, expected string
	&lt;cmdline&gt;:1:1-17
	During evaluation
$ jsonnet -e '{[a]: 1 for a in 2}'
RUNTIME ERROR: std.flatMap second param must be array / string, got number
	&lt;cmdline&gt;:1:1-20
	&lt;cmdline&gt;:1:1-20
	During evaluation
$ jsonnet -e '[a for a in 1]'
RUNTIME ERROR: std.flatMap second param must be array / string, got number
	&lt;cmdline&gt;:1:1-15
	During evaluation
```

FWIW, the C++ implementation does not have any of these problems. It
gives:
```
RUNTIME ERROR: In comprehension, can only iterate over array.
```

In the Go implementation comprehensions are desugared to a call to
std.flatMap which does accept a string in the "arr" parameter.

The fix: Desugar comprehensions to a call to a new hidden builtin which
only accepts arrays.
diff --git a/builtins.go b/builtins.go
@@ -345,6 +345,19 @@ func builtinFlatMap(i *interpreter, funcv, arrv value) (value, error) {
 	}
 }
 
+// builtinFlatMapArray is like builtinFlatMap, but only accepts array as the
+// arrv value. Desugared comprehensions contain a call to this function, rather
+// than builtinFlatMap, so that a better error message is printed when the
+// comprehension would iterate over a non-array.
+func builtinFlatMapArray(i *interpreter, funcv, arrv value) (value, error) {
+	switch arrv := arrv.(type) {
+	case *valueArray:
+		return builtinFlatMap(i, funcv, arrv)
+	default:
+		return nil, i.typeErrorSpecific(arrv, &valueArray{})
+	}
+}
+
 func joinArrays(i *interpreter, sep *valueArray, arr *valueArray) (value, error) {
 	result := make([]*cachedThunk, 0, arr.length())
 	first := true
@@ -2880,4 +2893,5 @@ var funcBuiltins = buildBuiltinMap([]builtin{
 
 	// internal
 	&unaryBuiltin{name: "$objectFlatMerge", function: builtinUglyObjectFlatMerge, params: ast.Identifiers{"x"}},
+	&binaryBuiltin{name: "$flatMapArray", function: builtinFlatMapArray, params: ast.Identifiers{"func", "arr"}},
 })
diff --git a/internal/program/desugarer.go b/internal/program/desugarer.go
@@ -184,7 +184,7 @@ func desugarForSpec(inside ast.Node, loc ast.LocationRange, forSpec *ast.ForSpec
 	if err != nil {
 		return nil, err
 	}
-	current := buildStdCall("flatMap", loc, function, forSpec.Expr)
+	current := buildStdCall("$flatMapArray", loc, function, forSpec.Expr)
 	if forSpec.Outer == nil {
 		return current, nil
 	}
diff --git a/linter/internal/types/stdlib.go b/linter/internal/types/stdlib.go
@@ -199,6 +199,7 @@ func prepareStdlib(g *typeGraph) {
 		"mod":              g.newSimpleFuncType(stringOrNumber, "a", "b"),
 		"native":           g.newSimpleFuncType(anyFunctionType, "x"),
 		"$objectFlatMerge": g.newSimpleFuncType(anyObjectType, "x"),
+		"$flatMapArray":    g.newSimpleFuncType(anyArrayType, "func", "arr"),
 
 		// Boolean
 
diff --git a/testdata/array_comp_try_iterate_over_empty_string.golden b/testdata/array_comp_try_iterate_over_empty_string.golden
@@ -0,0 +1,10 @@
+RUNTIME ERROR: Unexpected type string, expected array
+-------------------------------------------------
+	testdata/array_comp_try_iterate_over_empty_string:1:1-16	
+
+[a for a in '']
+
+-------------------------------------------------
+	During evaluation	
+
+
diff --git a/testdata/array_comp_try_iterate_over_empty_string.jsonnet b/testdata/array_comp_try_iterate_over_empty_string.jsonnet
@@ -0,0 +1 @@
+[a for a in '']
diff --git a/testdata/array_comp_try_iterate_over_empty_string.linter.golden b/testdata/array_comp_try_iterate_over_empty_string.linter.golden
diff --git a/testdata/array_comp_try_iterate_over_obj.golden b/testdata/array_comp_try_iterate_over_obj.golden
@@ -1,4 +1,4 @@
-RUNTIME ERROR: std.flatMap second param must be array / string, got object
+RUNTIME ERROR: Unexpected type object, expected array
 -------------------------------------------------
 	testdata/array_comp_try_iterate_over_obj:1:1-16	
 
diff --git a/testdata/array_comp_try_iterate_over_string.golden b/testdata/array_comp_try_iterate_over_string.golden
@@ -0,0 +1,10 @@
+RUNTIME ERROR: Unexpected type string, expected array
+-------------------------------------------------
+	testdata/array_comp_try_iterate_over_string:1:1-17	
+
+[a for a in 'b']
+
+-------------------------------------------------
+	During evaluation	
+
+
diff --git a/testdata/array_comp_try_iterate_over_string.jsonnet b/testdata/array_comp_try_iterate_over_string.jsonnet
@@ -0,0 +1 @@
+[a for a in 'b']
diff --git a/testdata/array_comp_try_iterate_over_string.linter.golden b/testdata/array_comp_try_iterate_over_string.linter.golden
diff --git a/testdata/object_comp_try_iterate_over_obj.golden b/testdata/object_comp_try_iterate_over_obj.golden
@@ -0,0 +1,15 @@
+RUNTIME ERROR: Unexpected type object, expected array
+-------------------------------------------------
+	testdata/object_comp_try_iterate_over_obj:1:1-23	
+
+{ [a]: 0 for a in {} }
+
+-------------------------------------------------
+	testdata/object_comp_try_iterate_over_obj:1:1-23	
+
+{ [a]: 0 for a in {} }
+
+-------------------------------------------------
+	During evaluation	
+
+
diff --git a/testdata/object_comp_try_iterate_over_obj.jsonnet b/testdata/object_comp_try_iterate_over_obj.jsonnet
@@ -0,0 +1 @@
+{ [a]: 0 for a in {} }
diff --git a/testdata/object_comp_try_iterate_over_obj.linter.golden b/testdata/object_comp_try_iterate_over_obj.linter.golden
diff --git a/testdata/object_comp_try_iterate_over_string.golden b/testdata/object_comp_try_iterate_over_string.golden
@@ -0,0 +1,15 @@
+RUNTIME ERROR: Unexpected type string, expected array
+-------------------------------------------------
+	testdata/object_comp_try_iterate_over_string:1:1-24	
+
+{ [a]: 0 for a in 'b' }
+
+-------------------------------------------------
+	testdata/object_comp_try_iterate_over_string:1:1-24	
+
+{ [a]: 0 for a in 'b' }
+
+-------------------------------------------------
+	During evaluation	
+
+
diff --git a/testdata/object_comp_try_iterate_over_string.jsonnet b/testdata/object_comp_try_iterate_over_string.jsonnet
@@ -0,0 +1 @@
+{ [a]: 0 for a in 'b' }
diff --git a/testdata/object_comp_try_iterate_over_string.linter.golden b/testdata/object_comp_try_iterate_over_string.linter.golden

Original file line number	Diff line number	Diff line change
`@@ -184,7 +184,7 @@ func desugarForSpec(inside ast.Node, loc ast.LocationRange, forSpec *ast.ForSpec`
`184`	`184`	`if err != nil {`
`185`	`185`	`return nil, err`
`186`	`186`	`}`
`187`		`- current := buildStdCall("flatMap", loc, function, forSpec.Expr)`
	`187`	`+ current := buildStdCall("$flatMapArray", loc, function, forSpec.Expr)`
`188`	`188`	`if forSpec.Outer == nil {`
`189`	`189`	`return current, nil`
`190`	`190`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-RUNTIME ERROR: std.flatMap second param must be array / string, got object`
	`1`	`+RUNTIME ERROR: Unexpected type object, expected array`
`2`	`2`	`-------------------------------------------------`
`3`	`3`	`testdata/array_comp_try_iterate_over_obj:1:1-16`
`4`	`4`