Protobuf API V2: compatibility fixes

go/protomodule/protomodule_map.go

@@ -60,11 +60,11 @@ func newProtoMapFromDict(mapKey protoreflect.FieldDescriptor, mapValue protorefl
 	out := &protoMap{
 		mapKey:   mapKey,
 		mapValue: mapValue,
-		dict:     d,
+		dict:     starlark.NewDict(d.Len()),
 	}
 
 	for _, item := range d.Items() {
-		err := out.typeCheck(item[0], item[1])
+		err := out.SetKey(item[0], item[1])
 		if err != nil {
 			return nil, err
 		}
@@ -145,19 +145,25 @@ func (m *protoMap) wrapUpdate() starlark.Value {
 }
 
 func (m *protoMap) SetKey(k, v starlark.Value) error {
-	err := m.typeCheck(k, v)
+	// Typecheck key
+	err := scalarTypeCheck(m.mapKey, k)
 	if err != nil {
 		return err
 	}
 
-	return m.dict.SetKey(k, v)
-}
+	// Pre 1.0 compatibility allowed maps to be constructed with None in proto2
+	// with the value treated as nil. protoreflect does not allow setting
+	// with a value of nil, so instead treat it as an unset
+	if fieldAllowsNone(m.mapValue) && v == starlark.None {
+		_, _, err = m.dict.Delete(k)
+		return err
+	}
 
-func (m *protoMap) typeCheck(k, v starlark.Value) error {
-	err := scalarTypeCheck(m.mapKey, k)
+	// Typecheck value
+	err = scalarTypeCheck(m.mapValue, v)
 	if err != nil {
 		return err
 	}
 
-	return scalarTypeCheck(m.mapValue, v)
+	return m.dict.SetKey(k, v)
 }

go/protomodule/protomodule_message.go

@@ -238,7 +238,7 @@ func (msg *protoMessage) SetField(name string, val starlark.Value) error {
 	}
 
 	// Allow using msg_field = None to unset a scalar message field
-	if fieldDesc.Kind() == protoreflect.MessageKind && val == starlark.None && !fieldDesc.IsList() && !fieldDesc.IsMap() {
+	if fieldAllowsNone(fieldDesc) && val == starlark.None {
 		delete(msg.fields, name)
 		return nil
 	}
@@ -367,3 +367,7 @@ func isFieldSet(v protoreflect.Value, fieldDesc protoreflect.FieldDescriptor) bo
 		return v.Interface() != fieldDesc.Default().Interface()
 	}
 }
+
+func fieldAllowsNone(fieldDesc protoreflect.FieldDescriptor) bool {
+	return (fieldDesc.Kind() == protoreflect.MessageKind && !fieldDesc.IsList() && !fieldDesc.IsMap()) || fieldDesc.Syntax() == protoreflect.Proto2
+}

go/protomodule/protomodule_message_test.go

@@ -379,10 +379,11 @@ func TestAttrValidation(t *testing.T) {
 		},
 
 		// Repeated and map fields can't be assigned `None`. Scalar fields can't be assigned `None`
+		// in proto3, but the error message is specialized.
 		{
 			name:    "none to scalar",
 			src:     `MessageV3(f_int32 = None)`,
-			wantErr: `TypeError: value None (type "NoneType") can't be assigned to type "int32".`,
+			wantErr: `TypeError: value None (type "NoneType") can't be assigned to type "int32" in proto3 mode.`,
 		},
 		{
 			name:    "none to string list",
@@ -966,3 +967,82 @@ func TestProtoMergeDiffTypes(t *testing.T) {
 		t.Errorf("expected error %q, got %q", errorMsg, err.Error())
 	}
 }
+
+// Pre 1.0 Skycfg allowed maps to be constructed with None values for proto2 (see protoMap.SetKey)
+func TestMapNoneCompatibility(t *testing.T) {
+	val, err := evalFunc(`
+def fun():
+    pb = proto.package("skycfg.test_proto")
+    msg = pb.MessageV2()
+    m = {
+	    "a": pb.MessageV2(),
+	    "b": pb.MessageV2(),
+	    "c": pb.MessageV2(),
+	    "d": None,
+    }
+    msg.map_submsg = m
+
+    m2 = msg.map_submsg
+    m2["b"] = None
+    m2.setdefault("e", None)
+    m2.update([("c", None)])
+
+    return msg
+`, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := mustProtoMessage(t, val).(*pb.MessageV2)
+
+	checkProtoEqual(t, &pb.MessageV2{
+		MapSubmsg: map[string]*pb.MessageV2{
+			"a": &pb.MessageV2{},
+		},
+	}, got)
+
+	// Confirm this only works for all in proto2, only message values in proto3
+	// This is an artifact of set to None being allow for scalar values in proto2
+	val, err = evalFunc(`
+def fun():
+    pb = proto.package("skycfg.test_proto")
+    msg = pb.MessageV2(
+	map_string = {
+	    "a": None
+        }
+    )
+    return msg
+`, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	val, err = evalFunc(`
+def fun():
+    pb = proto.package("skycfg.test_proto")
+    msg = pb.MessageV3(
+	map_string = {
+	    "a": None
+        }
+    )
+    return msg
+`, nil)
+	wantErr := fmt.Errorf(`TypeError: value None (type "NoneType") can't be assigned to type "string" in proto3 mode.`)
+	if !checkError(err, wantErr) {
+		t.Fatalf("eval: expected error %v, got %v", wantErr, err)
+	}
+}
+
+func TestUnsetProto2Fields(t *testing.T) {
+	// Proto v2 distinguishes between unset and set-to-empty.
+	msg, err := eval(`proto.package("skycfg.test_proto").MessageV2(
+                f_string = None,
+        )`, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := mustProtoMessage(t, msg)
+	want := &pb.MessageV2{
+		FString: nil,
+	}
+	checkProtoEqual(t, want, got)
+}

go/protomodule/type_conversions.go

@@ -327,8 +327,14 @@ func typeError(fieldDesc protoreflect.FieldDescriptor, val starlark.Value, scala
 		}
 	}
 
-	return fmt.Errorf("TypeError: value %s (type %q) can't be assigned to type %q.",
-		val.String(), val.Type(), expectedType,
+	// Add special error message for = None not allowed in proto3
+	proto3SpecialCase := ""
+	if scalar && val == starlark.None {
+		proto3SpecialCase = " in proto3 mode"
+	}
+
+	return fmt.Errorf("TypeError: value %s (type %q) can't be assigned to type %q%s.",
+		val.String(), val.Type(), expectedType, proto3SpecialCase,
 	)
 }
 

skycfg.go

@@ -197,14 +197,24 @@ func UnstablePredeclaredModules(r unstableProtoRegistryV2) starlark.StringDict {
 	return starlark.StringDict{
 		"fail":   assertmodule.Fail,
 		"hash":   hashmodule.NewModule(),
-		"json":   starlarkjson.Module,
+		"json":   newJsonModule(),
 		"proto":  UnstableProtoModule(r),
 		"struct": starlark.NewBuiltin("struct", starlarkstruct.Make),
 		"yaml":   newYamlModule(),
 		"url":    urlmodule.NewModule(),
 	}
 }
 
+func newJsonModule() starlark.Value {
+	module := starlarkjson.Module
+
+	// Aliases for compatibility with pre-v1.0 Skycfg API.
+	module.Members["marshal"] = module.Members["encode"]
+	module.Members["unmarshal"] = module.Members["decode"]
+
+	return module
+}
+
 func newYamlModule() starlark.Value {
 	module := yamlmodule.NewModule()