lib/json: allow encoding dicts with int keys

lib/json/json.go

@@ -155,19 +155,45 @@ func encode(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, k
 			// e.g. dict (must have string keys)
 			buf.WriteByte('{')
 			items := x.Items()
+			stringKeys := false
+			intKeys := false
 			for _, item := range items {
-				if _, ok := item[0].(starlark.String); !ok {
-					return fmt.Errorf("%s has %s key, want string", x.Type(), item[0].Type())
+				_, isString := item[0].(starlark.String)
+				_, isInt := item[0].(starlark.Int)
+				if !isString && !isInt {
+					return fmt.Errorf("%s has %s key, want string or int", x.Type(), item[0].Type())
 				}
+				stringKeys = stringKeys || isString
+				intKeys = intKeys || isInt
+			}
+			// Sort the keys. This is useful if they come from a dict, which presents
+			// its items unordered. However, we only sort if all the keys have the
+			// same type.
+			if !(stringKeys && intKeys) {
+				sort.Slice(items, func(i, j int) bool {
+					// Compare as strings if all the keys are strings, or as ints if all
+					// the keys are ints.
+					if stringKeys {
+						return items[i][0].(starlark.String) < items[j][0].(starlark.String)
+					}
+					cmp, err := items[i][0].(starlark.Int).Cmp(items[j][0], 0 /* depth */)
+					if err != nil {
+						panic(fmt.Errorf("unexpected failure to compare ints: %w", err))
+					}
+					return cmp == -1
+				})
 			}
-			sort.Slice(items, func(i, j int) bool {
-				return items[i][0].(starlark.String) < items[j][0].(starlark.String)
-			})
 			for i, item := range items {
 				if i > 0 {
 					buf.WriteByte(',')
 				}
-				k, _ := starlark.AsString(item[0])
+				key := item[0]
+				k, ok := starlark.AsString(key)
+				if !ok {
+					// If the key is not a string, it must be an int as per the checks
+					// above.
+					k = key.(starlark.Int).BigInt().String()
+				}
 				quote(k)
 				buf.WriteByte(':')
 				if err := emit(item[1]); err != nil {

starlark/testdata/json.star

@@ -22,20 +22,20 @@ assert.eq(json.encode((1, 2, 3)), "[1,2,3]")
 assert.eq(json.encode(range(3)), "[0,1,2]") # a built-in iterable
 assert.eq(json.encode(dict(x = 1, y = "two")), '{"x":1,"y":"two"}')
 assert.eq(json.encode(dict(y = "two", x = 1)), '{"x":1,"y":"two"}') # key, not insertion, order
+assert.eq(json.encode({1: "one", 2: "two"}), '{"1":"one","2":"two"}') # key, not insertion, order
 assert.eq(json.encode(struct(x = 1, y = "two")), '{"x":1,"y":"two"}')  # a user-defined HasAttrs
 assert.eq(json.encode("😹"[:1]), '"\\ufffd"') # invalid UTF-8 -> replacement char
 
 def encode_error(expr, error):
     assert.fails(lambda: json.encode(expr), error)
 
 encode_error(float("NaN"), "json.encode: cannot encode non-finite float nan")
-encode_error({1: "two"}, "dict has int key, want string")
+encode_error({(1,2): "two"}, "dict has tuple key, want string or int")
 encode_error(len, "cannot encode builtin_function_or_method as JSON")
 encode_error(struct(x=[1, {"x": len}]), # nested failure
              'in field .x: at list index 1: in dict key "x": cannot encode...')
 encode_error(struct(x=[1, {"x": len}]), # nested failure
              'in field .x: at list index 1: in dict key "x": cannot encode...')
-encode_error({1: 2}, 'dict has int key, want string')
 
 recursive_map = {}
 recursive_map["r"] = recursive_map