gen/enum: Add go.label to override item name

gen/enum.go

@@ -26,6 +26,11 @@ import (
 	"go.uber.org/thriftrw/compile"
 )
 
+const (
+	// GoLabel overrides name on the wire
+	GoLabel = "go.label"
+)
+
 // enumGenerator generates code to serialize and deserialize enums.
 type enumGenerator struct{}
 
@@ -98,7 +103,7 @@ func enum(g Generator, spec *compile.EnumSpec) error {
 			switch string(value) {
 			<- $enum := .Spec ->
 			<range .Spec.Items ->
-				case "<.Name>":
+				case "<enumItemWireName .>":
 					*<$v> = <enumItemName $enumName .>
 					return nil
 			<end ->
@@ -118,7 +123,7 @@ func enum(g Generator, spec *compile.EnumSpec) error {
 				switch int32(<$v>) {
 				<range .UniqueItems ->
 					case <.Value>:
-						return []byte("<.Name>"), nil
+						return []byte("<enumItemWireName .>"), nil
 				<end ->
 				}
 			<end ->
@@ -165,7 +170,7 @@ func enum(g Generator, spec *compile.EnumSpec) error {
 				switch <$w> {
 				<range .UniqueItems ->
 					case <.Value>:
-						return "<.Name>"
+						return "<enumItemWireName .>"
 				<end ->
 				}
 			<end ->
@@ -190,7 +195,7 @@ func enum(g Generator, spec *compile.EnumSpec) error {
 				switch int32(<$v>) {
 				<range .UniqueItems ->
 					case <.Value>:
-						return ([]byte)("\"<.Name>\""), nil
+						return ([]byte)("\"<enumItemWireName .>\""), nil
 				<end ->
 				}
 			<end ->
@@ -246,6 +251,7 @@ func enum(g Generator, spec *compile.EnumSpec) error {
 			UniqueItems: items,
 		},
 		TemplateFunc("enumItemName", enumItemName),
+		TemplateFunc("enumItemWireName", enumItemWireName),
 	)
 
 	return wrapGenerateError(spec.Name, err)
@@ -265,6 +271,17 @@ func enumItemName(enumName string, spec *compile.EnumItem) (string, error) {
 	return enumName + name, err
 }
 
+// enumItemWireName returns the actual name used for serialization/deserialization
+// default to EnumItem.Name, override by the value of GoLabel
+func enumItemWireName(spec *compile.EnumItem) string {
+	labelName := spec.Name
+	val, ok := spec.Annotations[GoLabel]
+	if ok && len(val) > 0 {
+		labelName = val
+	}
+	return labelName
+}
+
 // enumUniqueItems returns a subset of the given list of enum items where
 // there are no value collisions between items.
 func enumUniqueItems(items []compile.EnumItem) []compile.EnumItem {

gen/enum_test.go

@@ -399,3 +399,69 @@ func TestEnumAccessors(t *testing.T) {
 		})
 	})
 }
+
+func TestEnumLabelLegit(t *testing.T) {
+	tests := []struct {
+		enumItem  te.EnumWithLabel
+		wireValue string
+	}{
+		{te.EnumWithLabelUsername, "\"surname\""},
+		{te.EnumWithLabelPassword, "\"hashed_password\""},
+		{te.EnumWithLabelSalt, "\"salt\""},
+		{te.EnumWithLabelSugar, "\"sugar\""},
+	}
+	for _, tt := range tests {
+		t.Run("compare label:"+tt.wireValue, func(t *testing.T) {
+			// marshal
+			label := te.EnumWithLabel(tt.enumItem)
+			b, err := json.Marshal(label)
+			assert.NoError(t, err)
+			assert.Equal(t, tt.wireValue, string(b))
+
+			// unmarshal
+			var expectedLabel te.EnumWithLabel
+			err = json.Unmarshal(b, &expectedLabel)
+			assert.Equal(t, expectedLabel, label)
+
+			assertRoundTrip(
+				t, &tt.enumItem,
+				wire.NewValueI32(int32(tt.enumItem)),
+				"test roundtrip "+tt.wireValue,
+			)
+		})
+	}
+}
+
+func TestEnumLabelInvalid(t *testing.T) {
+
+	for _, tt := range []struct {
+		value te.EnumWithLabel
+		rep   string
+	}{
+		{-1, "-1"},
+		{1 << 10, "1024"},
+	} {
+		invalidEnumItem := te.EnumWithLabel(tt.value)
+		b, err := json.Marshal(invalidEnumItem)
+		assert.NoError(t, err)
+		assert.Equal(t, tt.rep, string(b))
+	}
+
+	for _, tt := range []struct {
+		errVal []byte
+		errMsg string
+	}{
+		{
+			[]byte("some-random-str"),
+			"invalid character 's' looking for beginning of value",
+		},
+		{
+			[]byte("\"; drop table users;\""),
+			"unknown enum value \"; drop table users;\" for \"EnumWithLabel\"",
+		},
+	} {
+		var expectedLabel te.EnumWithLabel
+		err := json.Unmarshal(tt.errVal, &expectedLabel)
+		assert.Equal(t, err.Error(), tt.errMsg)
+	}
+}

gen/internal/tests/thrift/enums.thrift

@@ -52,5 +52,13 @@ enum lowerCaseEnum {
     containing, lower_case, items
 }
 
+// EnumWithLabel use label name in serialization/deserialization
+enum EnumWithLabel {
+    username (go.label = "surname"),
+    password (go.label = "hashed_password"),
+    salt (go.label = ""),
+    sugar (go.label)
+}
+
 // collision with RecordType_Values() function.
 enum RecordType_Values { FOO, BAR }