diff --git a/references/cuegen/convert.go b/references/cuegen/convert.go index 98bd23774..0f4de1857 100644 --- a/references/cuegen/convert.go +++ b/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) diff --git a/references/cuegen/convert_test.go b/references/cuegen/convert_test.go index 8dbb678eb..e38e17b68 100644 --- a/references/cuegen/convert_test.go +++ b/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)) diff --git a/references/cuegen/option.go b/references/cuegen/option.go index 08e680b10..15eeec4da 100644 --- a/references/cuegen/option.go +++ b/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 + } +} diff --git a/references/cuegen/option_test.go b/references/cuegen/option_test.go index 9ac042cce..73b510b65 100644 --- a/references/cuegen/option_test.go +++ b/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) + } + } +} diff --git a/references/cuegen/testdata/valid.go b/references/cuegen/testdata/valid.go index bbfc4f222..8e235fcbc 100644 --- a/references/cuegen/testdata/valid.go +++ b/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"` +}