Merge pull request #2292 from gemmahou/gen-mapper

chore: Handle List of enum kind
GoogleCloudPlatform · Jul 16, 2024 · 746449b · 746449b
2 parents 727fcd4 + 0609b10
commit 746449b
Showing 1 changed file with 30 additions and 7 deletions.
diff --git a/dev/tools/controllerbuilder/pkg/codegen/mappergenerator.go b/dev/tools/controllerbuilder/pkg/codegen/mappergenerator.go
@@ -215,7 +215,7 @@ func (v *MapperGenerator) writeMapFunctionsForPair(out io.Writer, pair *typePair
 			}
 
 			if protoField.Cardinality() == protoreflect.Repeated {
-				useSliceFromProto := ""
+				useSliceFromProtoFunction := ""
 				useCustomMethod := false
 
 				switch protoField.Kind() {
@@ -225,12 +225,22 @@ func (v *MapperGenerator) writeMapFunctionsForPair(out io.Writer, pair *typePair
 					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "[]")
 
 					functionName := krmElemTypeName + "_FromProto"
-					useSliceFromProto = functionName
+					useSliceFromProtoFunction = functionName
 				case protoreflect.StringKind:
 					if krmField.Type != "[]string" {
 						useCustomMethod = true
 						// useSliceFromProto = fmt.Sprintf("%s_%s_FromProto", goTypeName, protoFieldName)
 					}
+				case protoreflect.EnumKind:
+					krmElemTypeName := krmField.Type
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "*")
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "[]")
+
+					functionName := "Enum_FromProto"
+					useSliceFromProtoFunction = fmt.Sprintf("%s(mapCtx, in.%s)",
+						functionName,
+						krmFieldName,
+					)
 
 				case
 					protoreflect.FloatKind,
@@ -241,7 +251,7 @@ func (v *MapperGenerator) writeMapFunctionsForPair(out io.Writer, pair *typePair
 					protoreflect.Uint32Kind,
 					protoreflect.Uint64Kind,
 					protoreflect.BytesKind:
-					useSliceFromProto = ""
+					useSliceFromProtoFunction = ""
 				default:
 					klog.Fatalf("unhandled kind %q for repeated field %v", protoField.Kind(), protoField)
 				}
@@ -251,19 +261,19 @@ func (v *MapperGenerator) writeMapFunctionsForPair(out io.Writer, pair *typePair
 					keyKind := entryMsg.Fields().ByName("key").Kind()
 					valueKind := entryMsg.Fields().ByName("value").Kind()
 					if keyKind == protoreflect.StringKind && valueKind == protoreflect.StringKind {
-						useSliceFromProto = ""
+						useSliceFromProtoFunction = ""
 					} else if keyKind == protoreflect.StringKind && valueKind == protoreflect.Int64Kind {
-						useSliceFromProto = ""
+						useSliceFromProtoFunction = ""
 					} else {
 						fmt.Fprintf(out, "// TODO: map type %v %v\n", keyKind, valueKind)
 					}
 				}
 
-				if useSliceFromProto != "" {
+				if useSliceFromProtoFunction != "" {
 					fmt.Fprintf(out, "\tout.%s = Slice_FromProto(mapCtx, in.%s, %s)\n",
 						krmFieldName,
 						krmFieldName,
-						useSliceFromProto,
+						useSliceFromProtoFunction,
 					)
 				} else if useCustomMethod {
 					methodName := fmt.Sprintf("%s_%s_FromProto", goTypeName, protoFieldName)
@@ -362,6 +372,19 @@ func (v *MapperGenerator) writeMapFunctionsForPair(out io.Writer, pair *typePair
 						//useSliceToProtoFunction = fmt.Sprintf("%s_%s_ToProto", goTypeName, protoFieldName)
 					}
 
+				case protoreflect.EnumKind:
+					krmElemTypeName := krmField.Type
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "*")
+					krmElemTypeName = strings.TrimPrefix(krmElemTypeName, "[]")
+
+					protoTypeName := "pb." + protoNameForEnum(protoField.Enum())
+					functionName := "Enum_ToProto"
+					useSliceToProtoFunction = fmt.Sprintf("%s[%s](mapCtx, in.%s)",
+						functionName,
+						protoTypeName,
+						krmFieldName,
+					)
+
 				case protoreflect.FloatKind,
 					protoreflect.DoubleKind,
 					protoreflect.BoolKind,