Feat: add type filter option (#5789)

Signed-off-by: iyear <ljyngup@gmail.com>

references/cuegen/convert.go

@@ -40,6 +40,10 @@ func (g *Generator) convertDecls(x *goast.GenDecl) (decls []cueast.Decl, _ error
 			continue
 		}
 
+		if g.opts.typeFilter != nil && !g.opts.typeFilter(typeSpec) {
+			continue
+		}
+
 		// only process struct
 		typ := g.pkg.TypesInfo.TypeOf(typeSpec.Name)
 

references/cuegen/convert_test.go

@@ -18,8 +18,10 @@ package cuegen
 
 import (
 	"bytes"
+	goast "go/ast"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -30,7 +32,14 @@ func TestConvert(t *testing.T) {
 	assert.NoError(t, err)
 
 	got := &bytes.Buffer{}
-	decls, err := g.Generate(WithAnyTypes("*k8s.io/apimachinery/pkg/apis/meta/v1/unstructured.Unstructured"))
+	decls, err := g.Generate(
+		WithAnyTypes("*k8s.io/apimachinery/pkg/apis/meta/v1/unstructured.Unstructured"),
+		WithTypeFilter(func(typ *goast.TypeSpec) bool {
+			if typ.Name == nil {
+				return true
+			}
+			return !strings.HasPrefix(typ.Name.Name, "TypeFilter")
+		}))
 	assert.NoError(t, err)
 	assert.NoError(t, g.Format(got, decls))
 

references/cuegen/option.go

@@ -16,9 +16,12 @@ limitations under the License.
 
 package cuegen
 
+import goast "go/ast"
+
 type options struct {
-	anyTypes map[string]struct{}
-	nullable bool
+	anyTypes   map[string]struct{}
+	nullable   bool
+	typeFilter func(typ *goast.TypeSpec) bool
 }
 
 var defaultOptions = &options{
@@ -26,7 +29,8 @@ var defaultOptions = &options{
 		"map[string]interface{}": {}, "map[string]any": {},
 		"interface{}": {}, "any": {},
 	},
-	nullable: false,
+	nullable:   false,
+	typeFilter: func(_ *goast.TypeSpec) bool { return true },
 }
 
 // Option is a function that configures generation options
@@ -51,3 +55,15 @@ func WithNullable() Option {
 		opts.nullable = true
 	}
 }
+
+// WithTypeFilter filters top struct types to be generated, and filter returns true to generate the type, otherwise false
+func WithTypeFilter(filter func(typ *goast.TypeSpec) bool) Option {
+	// return invalid option if filter is nil, so that it will not be applied
+	if filter == nil {
+		return nil
+	}
+
+	return func(opts *options) {
+		opts.typeFilter = filter
+	}
+}

references/cuegen/option_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package cuegen
 
 import (
+	goast "go/ast"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -73,3 +74,52 @@ func TestWithNullable(t *testing.T) {
 		assert.Equal(t, opts.nullable, tt.want, tt.name)
 	}
 }
+
+func TestWithTypeFilter(t *testing.T) {
+	tests := []struct {
+		name  string
+		opts  []Option
+		true  []string
+		false []string
+	}{
+		{
+			name:  "default",
+			opts:  nil,
+			true:  []string{"foo", "bar"},
+			false: []string{},
+		},
+		{
+			name: "nil",
+			opts: []Option{WithTypeFilter(nil)},
+			true: []string{"foo", "bar"},
+		},
+		{
+			name:  "single",
+			opts:  []Option{WithTypeFilter(func(typ *goast.TypeSpec) bool { return typ.Name.Name == "foo" })},
+			true:  []string{"foo"},
+			false: []string{"bar", "baz"},
+		},
+		{
+			name: "multiple",
+			opts: []Option{WithTypeFilter(func(typ *goast.TypeSpec) bool { return typ.Name.Name == "foo" }),
+				WithTypeFilter(func(typ *goast.TypeSpec) bool { return typ.Name.Name == "bar" })},
+			true:  []string{"bar"},
+			false: []string{"foo", "baz"},
+		},
+	}
+
+	for _, tt := range tests {
+		opts := options{typeFilter: func(_ *goast.TypeSpec) bool { return true }}
+		for _, opt := range tt.opts {
+			if opt != nil {
+				opt(&opts)
+			}
+		}
+		for _, typ := range tt.true {
+			assert.True(t, opts.typeFilter(&goast.TypeSpec{Name: &goast.Ident{Name: typ}}), tt.name)
+		}
+		for _, typ := range tt.false {
+			assert.False(t, opts.typeFilter(&goast.TypeSpec{Name: &goast.Ident{Name: typ}}), tt.name)
+		}
+	}
+}

references/cuegen/testdata/valid.go

@@ -321,3 +321,12 @@ type Skip struct {
 	Field3 string `json:"-"`
 	Field4 string `json:"field4"`
 }
+
+// TypeFilter should be ignored
+type TypeFilter http.Header
+
+// TypeFilterStruct should be ignored
+type TypeFilterStruct struct {
+	Field1 string `json:"field1"`
+	Field2 string `json:"field2"`
+}