protogen: add an option to rewrite import paths
This allows us to implement the import_prefix parameter in the v1
protoc-gen-go.
Drop support for import_prefix in protogen, and explicitly produce an
error if it is used in the v2 protoc-gen-go.
Change-Id: I66136b6b3affa3c0e9a93dc565619c90c42c0ecc
Reviewed-on: https://go-review.googlesource.com/138257
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
diff --git a/cmd/protoc-gen-go/main.go b/cmd/protoc-gen-go/main.go
index efd294b..8c72be8 100644
--- a/cmd/protoc-gen-go/main.go
+++ b/cmd/protoc-gen-go/main.go
@@ -15,15 +15,21 @@
)
func main() {
- var flags flag.FlagSet
- plugins := flags.String("plugins", "", "deprecated option")
- opts := &protogen.Options{
- ParamFunc: flags.Set,
- }
+ var (
+ flags flag.FlagSet
+ plugins = flags.String("plugins", "", "deprecated option")
+ importPrefix = flags.String("import_prefix", "", "deprecated option")
+ opts = &protogen.Options{
+ ParamFunc: flags.Set,
+ }
+ )
protogen.Run(opts, func(gen *protogen.Plugin) error {
if *plugins != "" {
return errors.New("protoc-gen-go: plugins are not supported; use 'protoc --go-grpc_out=...' to generate gRPC")
}
+ if *importPrefix != "" {
+ return errors.New("protoc-gen-go: import_prefix is not supported")
+ }
for _, f := range gen.Files {
if !f.Generate {
continue
diff --git a/protogen/protogen.go b/protogen/protogen.go
index 87e643c..80e7dc1 100644
--- a/protogen/protogen.go
+++ b/protogen/protogen.go
@@ -100,6 +100,7 @@
enumsByName map[protoreflect.FullName]*Enum
pathType pathType
genFiles []*GeneratedFile
+ opts *Options
err error
}
@@ -129,6 +130,11 @@
// if *value { ... }
// })
ParamFunc func(name, value string) error
+
+ // ImportRewriteFunc is called with the import path of each package
+ // imported by a generated file. It returns the import path to use
+ // for this package.
+ ImportRewriteFunc func(GoImportPath) GoImportPath
}
// New returns a new Plugin.
@@ -144,6 +150,7 @@
fileReg: protoregistry.NewFiles(),
messagesByName: make(map[protoreflect.FullName]*Message),
enumsByName: make(map[protoreflect.FullName]*Enum),
+ opts: opts,
}
packageNames := make(map[string]GoPackageName) // filename -> package name
@@ -158,8 +165,6 @@
switch param {
case "":
// Ignore.
- case "import_prefix":
- // TODO
case "import_path":
packageImportPath = GoImportPath(value)
case "paths":
@@ -712,6 +717,7 @@
// A GeneratedFile is a generated file.
type GeneratedFile struct {
+ gen *Plugin
filename string
goImportPath GoImportPath
buf bytes.Buffer
@@ -724,6 +730,7 @@
// and import path.
func (gen *Plugin) NewGeneratedFile(filename string, goImportPath GoImportPath) *GeneratedFile {
g := &GeneratedFile{
+ gen: gen,
filename: filename,
goImportPath: goImportPath,
packageNames: make(map[GoImportPath]GoPackageName),
@@ -876,8 +883,14 @@
importPaths = append(importPaths, string(importPath))
}
sort.Strings(importPaths)
+ rewriteImport := func(importPath string) string {
+ if f := g.gen.opts.ImportRewriteFunc; f != nil {
+ return string(f(GoImportPath(importPath)))
+ }
+ return importPath
+ }
for _, importPath := range importPaths {
- astutil.AddNamedImport(fset, file, string(g.packageNames[GoImportPath(importPath)]), importPath)
+ astutil.AddNamedImport(fset, file, string(g.packageNames[GoImportPath(importPath)]), rewriteImport(importPath))
}
for importPath := range g.manualImports {
if _, ok := g.packageNames[importPath]; ok {
diff --git a/protogen/protogen_test.go b/protogen/protogen_test.go
index 237226c..6d8e5af 100644
--- a/protogen/protogen_test.go
+++ b/protogen/protogen_test.go
@@ -307,6 +307,42 @@
}
}
+func TestImportRewrites(t *testing.T) {
+ gen, err := New(&pluginpb.CodeGeneratorRequest{}, &Options{
+ ImportRewriteFunc: func(i GoImportPath) GoImportPath {
+ return "prefix/" + i
+ },
+ })
+ if err != nil {
+ t.Fatal(err)
+ }
+ g := gen.NewGeneratedFile("foo.go", "golang.org/x/foo")
+ g.P("package foo")
+ g.P("var _ = ", GoIdent{GoName: "X", GoImportPath: "golang.org/x/bar"})
+ want := `package foo
+
+import bar "prefix/golang.org/x/bar"
+
+var _ = bar.X
+`
+ got, err := g.Content()
+ if err != nil {
+ t.Fatalf("g.Content() = %v", err)
+ }
+ if want != string(got) {
+ t.Fatalf(`want:
+==========
+%v
+==========
+
+got:
+==========
+%v
+==========`,
+ want, string(got))
+ }
+}
+
// makeRequest returns a CodeGeneratorRequest for the given protoc inputs.
//
// It does this by running protoc with the current binary as the protoc-gen-go