generating builders for enabled enum types (#31)

Author: dskatz

pkg/generators/builder/builder.go

@@ -187,7 +187,11 @@ func (b *BuilderPatternGenerator) generateSettersForType(sw *generator.SnippetWr
 		case m.Type == types.Bool:
 			sw.Do(setter.GenerateSetterForBool(m))
 		default:
-			sw.Do(setter.GenerateSetterForType(m))
+			if b.isTypePrimitiveEnabled(m) {
+				sw.Do(setter.GenerateSetterForTypeEnum(m, b.getEnabledPrimitiveType(m)))
+			} else {
+				sw.Do(setter.GenerateSetterForType(m))
+			}
 		}
 	}
 }
@@ -222,6 +226,11 @@ func (b *BuilderPatternGenerator) isTypePrimitiveEnabled(t types.Member) bool {
 	return exists
 }
 
+func (b *BuilderPatternGenerator) getEnabledPrimitiveType(t types.Member) *types.Type {
+	typeName := t.Name
+	return b.enabledTypes[typeName]
+}
+
 func (b *BuilderPatternGenerator) getWrapperType(t *types.Type) *types.Type {
 	typeName := t.Name.String()
 	if parent, ok := b.enabledTypes[typeName]; ok {

pkg/generators/snippets/setters.go

@@ -37,6 +37,23 @@ func (o $.pointer$$.type|raw$) $.funcName$(in $.memberType|raw$) $.pointer$$.typ
 	return raw, args
 }
 
+func (s *Setter) GenerateSetterForTypeEnum(member types.Member, argType *types.Type) (string, generator.Args) {
+	args := defaultGeneratorArgs(s.Root, s.pointerReceiver)
+	args["funcName"] = funcName(member)
+	args["memberAccessor"] = s.memberAccessor(member)
+	args["argType"] = argType
+	args["enumType"] = member.Type
+
+	raw := `// $.funcName$ is an autogenerated function
+func (o $.pointer$$.type|raw$) $.funcName$(in $.argType|raw$) $.pointer$$.type|raw$ {
+	o.$.memberAccessor$ = $.enumType|raw$(in)
+	return o
+}
+
+`
+	return raw, args
+}
+
 func (s *Setter) GenerateSetterForBool(member types.Member) (string, generator.Args) {
 	args := defaultGeneratorArgs(s.Root, s.pointerReceiver)
 	args["funcName"] = funcName(member)

pkg/generators/snippets/setters_test.go

@@ -466,6 +466,32 @@ func (o *SomeStruct) WithAlias(in Alias) *SomeStruct {
 	assert.Equal(t, want, b.String())
 }
 
+func TestGenerateSetterForStringEnum(t *testing.T) {
+	t.Parallel()
+
+	ctx, err := newTestGeneratorContext()
+	require.NoError(t, err)
+
+	someStruct := newTestType(t, "SomeStruct")
+	anEnumType := newTestType(t, "AnEnum")
+	root := someStruct
+	member := getMemberFromType(t, someStruct, "SomeStruct", "AnEnum")
+	parent := getMemberFromType(t, someStruct, "SomeStruct")
+	want := `// WithAnEnum is an autogenerated function
+func (o *SomeStruct) WithAnEnum(in AnEnum) *SomeStruct {
+	o.SomeStruct.AnEnum = b.AliasOfString(in)
+	return o
+}
+
+`
+	var b bytes.Buffer
+	setter := NewSetter(root, parent.Type, true)
+	sw := generator.NewSnippetWriter(&b, ctx, "$", "$")
+	sw.Do(setter.GenerateSetterForTypeEnum(member, anEnumType))
+	assert.NoError(t, sw.Error())
+	assert.Equal(t, want, b.String())
+}
+
 func TestMemberAccessor(t *testing.T) {
 	t.Parallel()
 

pkg/generators/snippets/testdata/a/a.go

@@ -19,6 +19,7 @@ type SomeStruct struct {
 	Bool               bool
 	PointerBool        *bool
 	Alias              *b.AliasOfString
+	AnEnum             b.AliasOfString
 }
 
 type AStruct struct {

pkg/generators/snippets/testdata/testdata.go

@@ -2,6 +2,7 @@ package testdata
 
 import (
 	"github.com/kanopy-platform/code-generator/pkg/generators/snippets/testdata/a"
+	"github.com/kanopy-platform/code-generator/pkg/generators/snippets/testdata/b"
 )
 
 type SomeStruct struct {
@@ -23,6 +24,7 @@ type ComplexStruct struct {
 }
 
 type Alias b.AliasOfString
+type AnEnum b.AliasOfString
 
 func (c *ComplexStruct) DeepCopyInto(in *ComplexStruct) {
 	// not impl test only