entproto: support enum values with special characters

entproto/README.md

@@ -430,6 +430,8 @@ The Proto Style Guide suggests that we use `CAPS_WITH_UNDERSCORES` for value nam
 - If no default value is defined for the enum, we generate a `<MessageName>_UNSPECIFIED = 0;` option on the enum and verify that no option received the 0 number in the enproto.Enum Options field.
 - If a default value is defined for the enum, we verify that it receives the 0 value on the Options field.
 
+Ent allows special characters in enum values. For such values, any special character is replaced by an underscore to preserve the `CAPS_WITH_UNDERSCORES` protobuf format.
+
 ## Edges
 
 Edges are annotated in the same way as fields: using `entproto.Field` annotation to specify the field number for the generated field. Unique relations are mapped to normal fields, non-unique relations are mapped to `repeated` fields.

entproto/adapter.go

@@ -441,7 +441,7 @@ func toProtoEnumDescriptor(fld *gen.Field) (*descriptorpb.EnumDescriptorProto, e
 		})
 	}
 	for _, opt := range fld.Enums {
-		n := strings.ToUpper(snake(opt.Value))
+		n := strings.ToUpper(snake(NormalizeEnumIdentifier(opt.Value)))
 		if !enumAnnotation.OmitFieldPrefix {
 			n = strings.ToUpper(snake(fld.Name)) + "_" + n
 		}

entproto/cmd/protoc-gen-entgrpc/main.go

@@ -125,6 +125,7 @@ func (g *serviceGenerator) generate() error {
 			"qualify": func(pkg, ident string) string {
 				return g.QualifiedGoIdent(protogen.GoImportPath(pkg).Ident(ident))
 			},
+			"protoIdentNormalize": entproto.NormalizeEnumIdentifier,
 			"statusErr": func(code, msg string) string {
 				return fmt.Sprintf("%s(%s, %q)",
 					g.QualifiedGoIdent(status.Ident("Error")),

entproto/cmd/protoc-gen-entgrpc/template/enums.tmpl

@@ -8,8 +8,13 @@
         {{ $entEnumIdent := entIdent $entLcase .PbStructField }}
         {{ $enumFieldPrefix := snake $enumType.GetName | upper | printf "%s_" }}
         {{ $omitPrefix := .EntField.Annotations.ProtoEnum.OmitFieldPrefix }}
-        func toProto{{ $pbEnumIdent.GoName }} (e {{ ident $entEnumIdent  }}) {{ ident $pbEnumIdent }} {
-            if v, ok := {{ $pbEnumIdent.GoName }}_value[{{ qualify "strings" "ToUpper" }}({{ if not $omitPrefix }}"{{ $enumFieldPrefix }}" +{{ end }} string(e))]; ok {
+        var protoIdentNormalizeRegexp{{ $pbEnumIdent.GoName }} = {{ qualify "regexp" "MustCompile" }}(`[^a-zA-Z0-9_]+`)
+        func protoIdentNormalize{{ $pbEnumIdent.GoName }}(e string) string {
+            return protoIdentNormalizeRegexp{{ $pbEnumIdent.GoName }}.ReplaceAllString(e, "_")
+        }
+
+        func toProto{{ $pbEnumIdent.GoName }} (e {{ ident $entEnumIdent }}) {{ ident $pbEnumIdent }} {
+            if v, ok := {{ $pbEnumIdent.GoName }}_value[{{ qualify "strings" "ToUpper" }}({{ if not $omitPrefix }}"{{ $enumFieldPrefix }}" + {{ end }}protoIdentNormalize{{ $pbEnumIdent.GoName }}(string(e)))]; ok {
                 return {{ $pbEnumIdent | ident }}(v)
             }
             return {{ $pbEnumIdent | ident }}(0)
@@ -19,7 +24,7 @@
             if v, ok := {{ $pbEnumIdent.GoName }}_name[int32(e)]; ok {
                 entVal := map[string]string{
                 {{- range .EntField.Enums }}
-                    "{{ if not $omitPrefix }}{{ $enumFieldPrefix }}{{ end }}{{ upper .Value }}": "{{ .Value }}",
+                    "{{ if not $omitPrefix }}{{ $enumFieldPrefix }}{{ end }}{{ protoIdentNormalize .Value }}": "{{ .Value }}",
                 {{- end }}
                 }[v]
                 return {{ ident $entEnumIdent }}(entVal)

entproto/enum.go

@@ -17,6 +17,8 @@ package entproto
 import (
 	"errors"
 	"fmt"
+	"regexp"
+	"strings"
 
 	"github.com/mitchellh/mapstructure"
 
@@ -29,6 +31,7 @@ const (
 
 var (
 	ErrEnumFieldsNotAnnotated = errors.New("entproto: all Enum options must be covered with an entproto.Enum annotation")
+	normalizeEnumIdent        = regexp.MustCompile(`[^a-zA-Z0-9_]+`)
 )
 
 type EnumOption func(*enum)
@@ -76,11 +79,18 @@ func (e *enum) Verify(fld *gen.Field) error {
 	if len(e.Options) != len(fld.Enums) {
 		return ErrEnumFieldsNotAnnotated
 	}
+	pbIdentifiers := make(map[string]struct{}, len(fld.Enums))
 	for _, opt := range fld.Enums {
 		if _, ok := e.Options[opt.Value]; !ok {
 			return fmt.Errorf("entproto: Enum option %s is not annotated with"+
 				" a pbfield number using entproto.Enum", opt.Name)
 		}
+		pbIdent := NormalizeEnumIdentifier(opt.Value)
+		if _, ok := pbIdentifiers[pbIdent]; ok {
+			return fmt.Errorf("entproto: Enum option %q produces conflicting pbfield"+
+				" name %q after normalization", opt.Name, pbIdent)
+		}
+		pbIdentifiers[pbIdent] = struct{}{}
 	}
 
 	// If default value is set on the pbfield, make sure it's option number is zero.
@@ -126,3 +136,9 @@ func extractEnumAnnotation(fld *gen.Field) (*enum, error) {
 
 	return &out, nil
 }
+
+// NormalizeEnumIdentifier normalizes the identifier of an enum pbfield
+// to match the Proto Style Guide.
+func NormalizeEnumIdentifier(s string) string {
+	return strings.ToUpper(normalizeEnumIdent.ReplaceAllString(s, "_"))
+}

entproto/internal/entprototest/adapter_test.go

@@ -123,6 +123,11 @@ func (suite *AdapterTestSuite) TestInvalidField() {
 	suite.EqualError(err, "unsupported field type \"TypeJSON\"")
 }
 
+func (suite *AdapterTestSuite) TestEnumWithConflictingValue() {
+	_, err := suite.adapter.GetFileDescriptor("EnumWithConflictingValue")
+	suite.EqualError(err, "entproto: Enum option \"EnumJpegAlt\" produces conflicting pbfield name \"IMAGE_JPEG\" after normalization")
+}
+
 func (suite *AdapterTestSuite) TestDuplicateNumber() {
 	_, err := suite.adapter.GetFileDescriptor("DuplicateNumberMessage")
 	suite.EqualError(err, "entproto: field 2 already defined on message \"DuplicateNumberMessage\"")
@@ -180,7 +185,7 @@ func (suite *AdapterTestSuite) TestEnumMessage() {
 	suite.NoError(err)
 
 	message := fd.FindMessage("entpb.MessageWithEnum")
-	suite.Len(message.GetFields(), 3)
+	suite.Len(message.GetFields(), 4)
 
 	// an enum field with defaults
 	enumField := message.FindFieldByName("enum_type")
@@ -202,6 +207,16 @@ func (suite *AdapterTestSuite) TestEnumMessage() {
 	suite.EqualValues(0, enumDesc.FindValueByName("ENUM_WITHOUT_DEFAULT_UNSPECIFIED").GetNumber())
 	suite.EqualValues(1, enumDesc.FindValueByName("ENUM_WITHOUT_DEFAULT_FIRST").GetNumber())
 	suite.EqualValues(2, enumDesc.FindValueByName("ENUM_WITHOUT_DEFAULT_SECOND").GetNumber())
+
+	// an enum field with special characters
+	enumField = message.FindFieldByName("enum_with_special_characters")
+	suite.EqualValues(4, enumField.GetNumber())
+	suite.EqualValues(descriptorpb.FieldDescriptorProto_TYPE_ENUM, enumField.GetType())
+	enumDesc = enumField.GetEnumType()
+	suite.EqualValues("entpb.MessageWithEnum.EnumWithSpecialCharacters", enumDesc.GetFullyQualifiedName())
+	suite.EqualValues(0, enumDesc.FindValueByName("ENUM_WITH_SPECIAL_CHARACTERS_UNSPECIFIED").GetNumber())
+	suite.EqualValues(1, enumDesc.FindValueByName("ENUM_WITH_SPECIAL_CHARACTERS_IMAGE_JPEG").GetNumber())
+	suite.EqualValues(2, enumDesc.FindValueByName("ENUM_WITH_SPECIAL_CHARACTERS_IMAGE_PNG").GetNumber())
 }
 
 func (suite *AdapterTestSuite) TestMessageWithId() {