Adds support for primitive alias types and enums (#28)

Author: dskatz

README.md

@@ -80,6 +80,29 @@ package mytypes
 // +kanopy:builder=package
 ```
 
+## Generate Enums
+
+An enum can be generated with the following argument. Enum constants are used in several upstream k8s packages.
+e.g.  
+
+```golang
+// +kanopy:builder=true,enum=*;CREATE;UPDATE;DELETE;CONNECT
+type OperationType admissionv1.OperationType
+```
+
+Which generates:
+```golang
+const OperationTypeCreate OperationType = "CREATE"
+```
+
+The `*` is a special value within k8s that indicates `All`.  The code generator will translate a `*` to all in the suffix of the constant name.
+
+```golang
+const OperationTypeAll OperationType = "*"
+```
+
+Enum arguments can be used for both upstream packages and internal packages.
+
 ## Definition of Terms
 
 | terms | definition |

pkg/generators/builder/builder.go

@@ -108,6 +108,11 @@ func (b *BuilderPatternGenerator) GenerateType(c *generator.Context, t *types.Ty
 
 	sw := generator.NewSnippetWriter(w, c, "$", "$")
 
+	if t.IsPrimitive() {
+		sw.Do(snippets.GenerateEnumSetter(t, tags.GetEnumOptions(t)))
+		return sw.Error()
+	}
+
 	if hasObjectMetaEmbedded(t) {
 		parentTypeOfObjectMeta := getParentOfEmbeddedType(t, ObjectMeta)
 		objectMetaType := getMemberTypeFromType(parentTypeOfObjectMeta, ObjectMeta)
@@ -172,6 +177,10 @@ func (b *BuilderPatternGenerator) generateSettersForType(sw *generator.SnippetWr
 				if b.isTypeEnabled(pointerType) {
 					sw.Do(setter.GenerateSetterForEmbeddedPointer(m, b.getWrapperType(pointerType)))
 				}
+			case types.Alias:
+				if b.isTypePrimitiveEnabled(m) {
+					sw.Do(setter.GenerateSetterForAliasPointerPrimitive(m, m.Type))
+				}
 			default:
 				sw.Do(setter.GenerateSetterForType(m))
 			}
@@ -207,6 +216,12 @@ func (b *BuilderPatternGenerator) isTypeEnabled(t *types.Type) bool {
 	return exists
 }
 
+func (b *BuilderPatternGenerator) isTypePrimitiveEnabled(t types.Member) bool {
+	typeName := t.Name
+	_, exists := b.enabledTypes[typeName]
+	return exists
+}
+
 func (b *BuilderPatternGenerator) getWrapperType(t *types.Type) *types.Type {
 	typeName := t.Name.String()
 	if parent, ok := b.enabledTypes[typeName]; ok {
@@ -257,16 +272,28 @@ func (b *BuilderPatternGenerator) Filter(c *generator.Context, t *types.Type) bo
 		return false
 	}
 
+	b.enablePrimitiveType(t)
+
 	for _, m := range t.Members {
-		if m.Embedded {
-			childType := m.Type.Name.String()
-			b.enabledTypes[childType] = t
-		}
+		b.enableEmbededMemberType(m, t)
 	}
 
 	return true
 }
 
+func (b *BuilderPatternGenerator) enableEmbededMemberType(m types.Member, t *types.Type) {
+	if m.Embedded {
+		childType := m.Type.Name.String()
+		b.enabledTypes[childType] = t
+	}
+}
+
+func (b *BuilderPatternGenerator) enablePrimitiveType(t *types.Type) {
+	if t.IsPrimitive() {
+		b.enabledTypes[t.Name.Name] = t
+	}
+}
+
 func (b *BuilderPatternGenerator) Namers(c *generator.Context) namer.NameSystems {
 	return namer.NameSystems{
 		"raw": namer.NewRawNamer(b.pkgToBuild.Path, b.imports),

pkg/generators/builder/builder_test.go

@@ -126,6 +126,30 @@ func TestBuilderPattern_NonObjectMetaGeneratesSnippets(t *testing.T) {
 	assert.NotContains(t, buf.String(), "DeepCopyInto")
 }
 
+func TestBuilderAliasPrimitiveType(t *testing.T) {
+	b := &BuilderPatternGeneratorFactory{}
+	pkg, typeToGenerate := newTestGeneratorType(t, "d", "DPolicyRule")
+	_, aliasToGenerate := newTestGeneratorType(t, "d", "AliasType")
+	g := b.NewBuilder(pkg)
+	buf := &bytes.Buffer{}
+	c := newGeneratorContext(g)
+	assert.True(t, g.Filter(c, typeToGenerate))
+	assert.True(t, g.Filter(c, aliasToGenerate))
+	assert.NoError(t, g.GenerateType(c, typeToGenerate, buf))
+	assert.Contains(t, buf.String(), "func (o *DPolicyRule) WithAliasType(in AliasType) *DPolicyRule")
+}
+
+func TestBuilderAliasPrimitiveTypeNotGenerated(t *testing.T) {
+	b := &BuilderPatternGeneratorFactory{}
+	pkg, typeToGenerate := newTestGeneratorType(t, "d", "DPolicyRule")
+	g := b.NewBuilder(pkg)
+	buf := &bytes.Buffer{}
+	c := newGeneratorContext(g)
+	assert.True(t, g.Filter(c, typeToGenerate))
+	assert.NoError(t, g.GenerateType(c, typeToGenerate, buf))
+	assert.NotContains(t, buf.String(), "func (o *DPolicyRule) WithAliasType(in AliasType) *DPolicyRule")
+}
+
 func TestBuilderPattern_GenerateSettersForType(t *testing.T) {
 	b := &BuilderPatternGeneratorFactory{}
 	pkg, typeToGenerate := newTestGeneratorType(t, "c", "CDeployment")

pkg/generators/builder/testdata/d/d.go

@@ -8,3 +8,6 @@ import (
 type DPolicyRule struct {
 	e.MockPolicyRule
 }
+
+// +kanopy:builder=true
+type AliasType e.AliasToString

pkg/generators/builder/testdata/d/e/e.go

@@ -3,4 +3,7 @@ package e
 type MockPolicyRule struct {
 	Verbs      []string
 	ListOfInts []int
+	AliasType  *AliasToString
 }
+
+type AliasToString string

pkg/generators/snippets/enum.go

@@ -0,0 +1,38 @@
+package snippets
+
+import (
+	"fmt"
+
+	"golang.org/x/text/cases"
+	"golang.org/x/text/language"
+	"k8s.io/gengo/generator"
+	"k8s.io/gengo/types"
+)
+
+const allValue = "*"
+const allSuffix = "All"
+
+func GenerateEnumSetter(inputType *types.Type, enumOptions []string) (string, generator.Args) {
+	args := generator.Args{
+		"type": inputType,
+		"name": inputType.Name.Name,
+	}
+
+	raw := ""
+
+	for _, val := range enumOptions {
+		raw += fmt.Sprintf(`const $.name$%s $.name$ = "%s"`, toSuffix(val), val)
+		raw += "\n"
+	}
+
+	return raw, args
+}
+
+func toSuffix(v string) string {
+	suffix := v
+	if suffix == allValue {
+		suffix = allSuffix
+	}
+	caser := cases.Title(language.English)
+	return caser.String(suffix)
+}

pkg/generators/snippets/enum_test.go

@@ -0,0 +1,60 @@
+package snippets
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"k8s.io/gengo/generator"
+	"k8s.io/gengo/types"
+)
+
+func TestGenerateEnumSetter(t *testing.T) {
+	t.Parallel()
+
+	ctx, err := newTestGeneratorContext()
+	assert.NoError(t, err)
+
+	tests := []struct {
+		description string
+		enumVals    []string
+		want        string
+	}{
+		{
+			description: "single value",
+			enumVals:    []string{"val1"},
+			want:        "const MyEnumVal1 MyEnum = \"val1\"\n",
+		},
+		{
+			description: "multi value",
+			enumVals:    []string{"val1", "val2"},
+			want:        "const MyEnumVal1 MyEnum = \"val1\"\nconst MyEnumVal2 MyEnum = \"val2\"\n",
+		},
+		{
+			description: "all caps",
+			enumVals:    []string{"VAL"},
+			want:        "const MyEnumVal MyEnum = \"VAL\"\n",
+		},
+	}
+
+	tt := enumTestType()
+
+	for _, test := range tests {
+		var b bytes.Buffer
+		sw := generator.NewSnippetWriter(&b, ctx, "$", "$")
+		sw.Do(GenerateEnumSetter(&tt, test.enumVals))
+		assert.NoError(t, sw.Error(), test.description)
+		assert.Equal(t, test.want, b.String(), test.description)
+	}
+}
+
+func enumTestType() types.Type {
+	tt := types.Type{
+		Name: types.Name{
+			Package: "./",
+			Name:    "MyEnum",
+			Path:    "",
+		},
+	}
+	return tt
+}

pkg/generators/snippets/setters.go

@@ -231,6 +231,23 @@ func (o $.pointer$$.type|raw$) $.funcName$(in *$.inputType|raw$) $.pointer$$.typ
 	return raw, args
 }
 
+func (s *Setter) GenerateSetterForAliasPointerPrimitive(member types.Member, inputType *types.Type) (string, generator.Args) {
+	args := defaultGeneratorArgs(s.Root, s.pointerReceiver)
+	args["funcName"] = funcName(member)
+	args["memberAccessor"] = member.Name
+	args["inputType"] = inputType.Elem
+
+	raw := `// $.funcName$ is an autogenerated function
+func (o $.pointer$$.type|raw$) $.funcName$(in $.memberAccessor$) $.pointer$$.type|raw$ {
+	p := $.inputType|raw$(in)
+	o.$.memberAccessor$ = &p
+	return o
+}
+
+`
+	return raw, args
+}
+
 func funcName(m types.Member) string {
 	verb := "With"
 

pkg/generators/snippets/setters_test.go

@@ -441,6 +441,31 @@ func TestFuncName(t *testing.T) {
 	}
 }
 
+func TestGenerateSetterForAliasPointerPrimitive(t *testing.T) {
+	t.Parallel()
+
+	ctx, err := newTestGeneratorContext()
+	require.NoError(t, err)
+
+	someStruct := newTestType(t, "SomeStruct")
+	root := someStruct
+	member := getMemberFromType(t, someStruct, "SomeStruct", "Alias")
+	want := `// WithAlias is an autogenerated function
+func (o *SomeStruct) WithAlias(in Alias) *SomeStruct {
+	p := b.AliasOfString(in)
+	o.Alias = &p
+	return o
+}
+
+`
+	var b bytes.Buffer
+	setter := NewSetter(root, member.Type, true)
+	sw := generator.NewSnippetWriter(&b, ctx, "$", "$")
+	sw.Do(setter.GenerateSetterForAliasPointerPrimitive(member, member.Type))
+	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

@@ -18,6 +18,7 @@ type SomeStruct struct {
 	MapStringByteSlice map[string][]byte
 	Bool               bool
 	PointerBool        *bool
+	Alias              *b.AliasOfString
 }
 
 type AStruct struct {