Merge pull request #22741 from wojtek-t/refactor_protobuf_generator

cmd/libs/go2idl/go-to-protobuf/protobuf/cmd.go

@@ -229,7 +229,7 @@ func Run(g *Generator) {
 
 		// alter the generated protobuf file to remove the generated types (but leave the serializers) and rewrite the
 		// package statement to match the desired package name
-		if err := RewriteGeneratedGogoProtobufFile(outputPath, p.GoPackageName(), p.ExtractGeneratedType, buf.Bytes()); err != nil {
+		if err := RewriteGeneratedGogoProtobufFile(outputPath, p.ExtractGeneratedType, buf.Bytes()); err != nil {
 			log.Fatalf("Unable to rewrite generated %s: %v", outputPath, err)
 		}
 

cmd/libs/go2idl/go-to-protobuf/protobuf/generator.go

@@ -347,8 +347,6 @@ type protoField struct {
 	Extras   map[string]string
 
 	CommentLines string
-
-	OptionalSet bool
 }
 
 var (
@@ -473,6 +471,7 @@ func protobufTagToField(tag string, field *protoField, m types.Member, t *types.
 	}
 	field.Tag = protoTag
 	// TODO: we are converting a Protobuf type back into an internal type, which is questionable
+	// TODO; Allow a way to unambiguously represent a type into two systems at the same time, like Go and Protobuf.
 	if last := strings.LastIndex(parts[0], "."); last != -1 {
 		prefix := parts[0][:last]
 		field.Type = &types.Type{
@@ -498,7 +497,6 @@ func protobufTagToField(tag string, field *protoField, m types.Member, t *types.
 			}
 		}
 	}
-	field.OptionalSet = true
 
 	protoExtra := make(map[string]string)
 	for i, extra := range parts[3:] {
@@ -511,6 +509,7 @@ func protobufTagToField(tag string, field *protoField, m types.Member, t *types.
 			parts[0] = fmt.Sprintf("(gogoproto.%s)", parts[0])
 			protoExtra[parts[0]] = parts[1]
 		}
+		// TODO: Should we parse castkey and castvalue too?
 	}
 
 	field.Extras = protoExtra
@@ -710,18 +709,6 @@ func (ft protoIDLFileType) assemble(w io.Writer, f *generator.File) {
 	w.Write(f.Body.Bytes())
 }
 
-func isPackable(t *types.Type) bool {
-	if t.Kind != typesKindProtobuf {
-		return false
-	}
-	switch t.Name.Name {
-	case "int32", "int64", "varint":
-		return true
-	default:
-		return false
-	}
-}
-
 func isPrivateGoName(name string) bool {
 	if len(name) == 0 {
 		return true

cmd/libs/go2idl/go-to-protobuf/protobuf/import_tracker.go

@@ -27,6 +27,8 @@ import (
 //
 // TODO: pay attention to the package name (instead of renaming every package).
 // TODO: Figure out the best way to make names for packages that collide.
+//
+// TODO: Try to merge this into generator.ImportTracker (or refactor common parts).
 type ImportTracker struct {
 	pathToName map[string]string
 	// forbidden names are in here. (e.g. "go" is a directory in which

cmd/libs/go2idl/go-to-protobuf/protobuf/namer.go

@@ -90,24 +90,6 @@ func (n *protobufNamer) GoNameToProtoName(name types.Name) types.Name {
 	return types.Name{Name: name.Name}
 }
 
-func (n *protobufNamer) protoNameForTypeName(name types.Name) string {
-	packageName := name.Package
-	if len(name.Package) != 0 {
-		if p, ok := n.packagesByPath[packageName]; ok {
-			packageName = p.Name()
-		} else {
-			packageName = protoSafePackage(packageName)
-		}
-	}
-	if len(name.Name) == 0 {
-		return packageName
-	}
-	if len(packageName) > 0 {
-		return packageName + "." + name.Name
-	}
-	return name.Name
-}
-
 func protoSafePackage(name string) string {
 	return strings.Replace(name, "/", ".", -1)
 }

cmd/libs/go2idl/go-to-protobuf/protobuf/package.go

@@ -32,44 +32,37 @@ import (
 
 func newProtobufPackage(packagePath, packageName string, generateAll bool, omitFieldTypes map[types.Name]struct{}) *protobufPackage {
 	pkg := &protobufPackage{
-		// The protobuf package name (foo.bar.baz)
-		PackageName: packageName,
-		// A path segment relative to the GOPATH root (foo/bar/baz)
-		PackagePath:    packagePath,
-		GenerateAll:    generateAll,
-		OmitFieldTypes: omitFieldTypes,
-		HeaderText: []byte(
-			`
-// This file was autogenerated by the command:
-// $ ` + os.Args[0] + `
-// Do not edit it manually!
+		DefaultPackage: generator.DefaultPackage{
+			// The protobuf package name (foo.bar.baz)
+			PackageName: packageName,
+			// A path segment relative to the GOPATH root (foo/bar/baz)
+			PackagePath: packagePath,
+			HeaderText: []byte(
+				`
+// This file was autogenerated by go-to-protobuf. Do not edit it manually!
 
 `),
-		PackageDocumentation: []byte(fmt.Sprintf(
-			`// Package %s is an autogenerated protobuf IDL.
+			PackageDocumentation: []byte(fmt.Sprintf(
+				`// Package %s is an autogenerated protobuf IDL.
 `, packageName)),
+		},
+		GenerateAll:    generateAll,
+		OmitFieldTypes: omitFieldTypes,
 	}
+	pkg.FilterFunc = pkg.filterFunc
+	pkg.GeneratorFunc = pkg.generatorFunc
 	return pkg
 }
 
 // protobufPackage contains the protobuf implementation of Package.
 type protobufPackage struct {
-	// Short name of package, used in the "package xxxx" line.
-	PackageName string
-	// Import path of the package, and the location on disk of the package.
-	PackagePath string
+	generator.DefaultPackage
+
 	// If true, generate protobuf serializations for all public types.
 	// If false, only generate protobuf serializations for structs that
 	// request serialization.
 	GenerateAll bool
 
-	// Emitted at the top of every file.
-	HeaderText []byte
-
-	// Emitted only for a "doc.go" file; appended to the HeaderText for
-	// that file.
-	PackageDocumentation []byte
-
 	// A list of types to filter to; if not specified all types will be included.
 	FilterTypes map[types.Name]struct{}
 
@@ -106,10 +99,7 @@ func (p *protobufPackage) ProtoTypeName() types.Name {
 	}
 }
 
-func (p *protobufPackage) Name() string { return p.PackageName }
-func (p *protobufPackage) Path() string { return p.PackagePath }
-
-func (p *protobufPackage) Filter(c *generator.Context, t *types.Type) bool {
+func (p *protobufPackage) filterFunc(c *generator.Context, t *types.Type) bool {
 	switch t.Kind {
 	case types.Func, types.Chan:
 		return false
@@ -175,7 +165,7 @@ func (p *protobufPackage) ExtractGeneratedType(t *ast.TypeSpec) bool {
 	return true
 }
 
-func (p *protobufPackage) Generators(c *generator.Context) []generator.Generator {
+func (p *protobufPackage) generatorFunc(c *generator.Context) []generator.Generator {
 	generators := []generator.Generator{}
 
 	p.Imports.AddNullable()
@@ -194,13 +184,6 @@ func (p *protobufPackage) Generators(c *generator.Context) []generator.Generator
 	return generators
 }
 
-func (p *protobufPackage) Header(filename string) []byte {
-	if filename == "doc.go" {
-		return append(p.HeaderText, p.PackageDocumentation...)
-	}
-	return p.HeaderText
-}
-
 func (p *protobufPackage) GoPackageName() string {
 	return filepath.Base(p.PackagePath)
 }

cmd/libs/go2idl/go-to-protobuf/protobuf/parser.go

@@ -33,11 +33,7 @@ import (
 	customreflect "k8s.io/kubernetes/third_party/golang/reflect"
 )
 
-// ExtractFunc extracts information from the provided TypeSpec and returns true if the type should be
-// removed from the destination file.
-type ExtractFunc func(*ast.TypeSpec) bool
-
-func RewriteGeneratedGogoProtobufFile(name string, packageName string, extractFn ExtractFunc, header []byte) error {
+func rewriteFile(name string, header []byte, rewriteFn func(*token.FileSet, *ast.File) error) error {
 	fset := token.NewFileSet()
 	src, err := ioutil.ReadFile(name)
 	if err != nil {
@@ -47,19 +43,10 @@ func RewriteGeneratedGogoProtobufFile(name string, packageName string, extractFn
 	if err != nil {
 		return err
 	}
-	cmap := ast.NewCommentMap(fset, file, file.Comments)
 
-	// remove types that are already declared
-	decls := []ast.Decl{}
-	for _, d := range file.Decls {
-		if !dropExistingTypeDeclarations(d, extractFn) {
-			decls = append(decls, d)
-		}
+	if err := rewriteFn(fset, file); err != nil {
+		return err
 	}
-	file.Decls = decls
-
-	// remove unmapped comments
-	file.Comments = cmap.Filter(file).Comments()
 
 	b := &bytes.Buffer{}
 	b.Write(header)
@@ -83,6 +70,29 @@ func RewriteGeneratedGogoProtobufFile(name string, packageName string, extractFn
 	return f.Close()
 }
 
+// ExtractFunc extracts information from the provided TypeSpec and returns true if the type should be
+// removed from the destination file.
+type ExtractFunc func(*ast.TypeSpec) bool
+
+func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, header []byte) error {
+	return rewriteFile(name, header, func(fset *token.FileSet, file *ast.File) error {
+		cmap := ast.NewCommentMap(fset, file, file.Comments)
+
+		// remove types that are already declared
+		decls := []ast.Decl{}
+		for _, d := range file.Decls {
+			if !dropExistingTypeDeclarations(d, extractFn) {
+				decls = append(decls, d)
+			}
+		}
+		file.Decls = decls
+
+		// remove unmapped comments
+		file.Comments = cmap.Filter(file).Comments()
+		return nil
+	})
+}
+
 func dropExistingTypeDeclarations(decl ast.Decl, extractFn ExtractFunc) bool {
 	switch t := decl.(type) {
 	case *ast.GenDecl:
@@ -108,52 +118,25 @@ func dropExistingTypeDeclarations(decl ast.Decl, extractFn ExtractFunc) bool {
 }
 
 func RewriteTypesWithProtobufStructTags(name string, structTags map[string]map[string]string) error {
-	fset := token.NewFileSet()
-	src, err := ioutil.ReadFile(name)
-	if err != nil {
-		return err
-	}
-	file, err := parser.ParseFile(fset, name, src, parser.DeclarationErrors|parser.ParseComments)
-	if err != nil {
-		return err
-	}
-
-	allErrs := []error{}
+	return rewriteFile(name, []byte{}, func(fset *token.FileSet, file *ast.File) error {
+		allErrs := []error{}
 
-	// set any new struct tags
-	for _, d := range file.Decls {
-		if errs := updateStructTags(d, structTags, []string{"protobuf"}); len(errs) > 0 {
-			allErrs = append(allErrs, errs...)
+		// set any new struct tags
+		for _, d := range file.Decls {
+			if errs := updateStructTags(d, structTags, []string{"protobuf"}); len(errs) > 0 {
+				allErrs = append(allErrs, errs...)
+			}
 		}
-	}
 
-	if len(allErrs) > 0 {
-		var s string
-		for _, err := range allErrs {
-			s += err.Error() + "\n"
+		if len(allErrs) > 0 {
+			var s string
+			for _, err := range allErrs {
+				s += err.Error() + "\n"
+			}
+			return errors.New(s)
 		}
-		return errors.New(s)
-	}
-
-	b := &bytes.Buffer{}
-	if err := printer.Fprint(b, fset, file); err != nil {
-		return err
-	}
-
-	body, err := format.Source(b.Bytes())
-	if err != nil {
-		return fmt.Errorf("%s\n---\nunable to format %q: %v", b, name, err)
-	}
-
-	f, err := os.OpenFile(name, os.O_WRONLY|os.O_TRUNC, 0644)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-	if _, err := f.Write(body); err != nil {
-		return err
-	}
-	return f.Close()
+		return nil
+	})
 }
 
 func updateStructTags(decl ast.Decl, structTags map[string]map[string]string, toCopy []string) []error {