add support for import map and override for GoNames.GoPackageForFile;…

… adds test for that method (#25)

plugins/names.go

@@ -19,6 +19,12 @@ import (
 //
 // GoNames is not thread-safe.
 type GoNames struct {
+	// A user-provided map of proto file names to the Go import package where
+	// that file's code is generated. These mappings can be specified via
+	// "M<protofile>=<gopkg>" args to the "--go_out" protoc argument. Other
+	// plugins that generate Go code may support the option, too.
+	ImportMap map[string]string
+
 	// cache of descriptor to TypeName
 	descTypes map[typeKey]gopoet.TypeName
 	// cache of descriptor to names
@@ -92,15 +98,33 @@ func (n *GoNames) OutputFilenameFor(fd *desc.FileDescriptor, suffix string) stri
 	return path.Join(pkg.ImportPath, name)
 }
 
-// GoPackageForFile returns the Go package for the given file descriptor.
+// GoPackageForFile returns the Go package for the given file descriptor. This will use
+// the file's "go_package" option if it has one, but that can be overridden if the user
+// has supplied an entry in n.ImportMap.
 func (n *GoNames) GoPackageForFile(fd *desc.FileDescriptor) gopoet.Package {
+	return n.GoPackageForFileWithOverride(fd, "")
+}
+
+// GoPackageForFileWithOverride returns the Go package for the given file descriptor,
+// but uses the given string as if it were the "go_package" option value.
+func (n *GoNames) GoPackageForFileWithOverride(fd *desc.FileDescriptor, goPackage string) gopoet.Package {
 	if pkg, ok := n.pkgNames[fd]; ok {
 		return pkg
 	}
 
-	fileName, protoPackage, goOption := fd.GetName(), fd.GetPackage(), fd.GetFileOptions().GetGoPackage()
+	// if not supplied: get go_package option from file, but allow it to
+	// be overridden by user-supplied import map
+	if goPackage == "" {
+		var ok bool
+		goPackage, ok = n.ImportMap[fd.GetName()]
+		if !ok {
+			goPackage = fd.GetFileOptions().GetGoPackage()
+		}
+	}
+
+	fileName, protoPackage := fd.GetName(), fd.GetPackage()
 	var pkgPath, pkgName string
-	if goOption == "" {
+	if goPackage == "" {
 		pkgPath = path.Dir(fileName)
 		if protoPackage == "" {
 			n := path.Base(fileName)
@@ -113,9 +137,8 @@ func (n *GoNames) GoPackageForFile(fd *desc.FileDescriptor) gopoet.Package {
 		} else {
 			pkgName = protoPackage
 		}
-		pkgName = sanitize(pkgName)
 	} else {
-		parts := strings.Split(goOption, ";")
+		parts := strings.Split(goPackage, ";")
 		if len(parts) > 1 {
 			pkgPath = parts[0]
 			pkgName = parts[1]
@@ -128,6 +151,7 @@ func (n *GoNames) GoPackageForFile(fd *desc.FileDescriptor) gopoet.Package {
 			}
 		}
 	}
+	pkgName = sanitize(pkgName)
 
 	pkg := gopoet.Package{ImportPath: pkgPath, Name: pkgName}
 	if n.pkgNames == nil {

plugins/names_test.go

@@ -0,0 +1,150 @@
+package plugins
+
+import (
+	"testing"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/golang/protobuf/protoc-gen-go/descriptor"
+	"github.com/jhump/gopoet"
+	"github.com/jhump/protoreflect/desc"
+	"github.com/jhump/protoreflect/desc/builder"
+)
+
+func mustBuildFile(f *builder.FileBuilder) *desc.FileDescriptor {
+	fd, err := f.Build()
+	if err != nil {
+		panic(err)
+	}
+	return fd
+}
+
+func TestGoPackageForFile(t *testing.T) {
+	testCases := []struct {
+		fd             *desc.FileDescriptor
+		importMap      map[string]string
+		override       string
+		expectedResult gopoet.Package
+	}{
+		{
+			fd:             mustBuildFile(builder.NewFile("foo/test.proto")),
+			expectedResult: gopoet.Package{ImportPath: "foo", Name: "test"},
+		},
+		{
+			fd: mustBuildFile(builder.NewFile("foo/test.proto").
+				SetPackageName("foo.bar")),
+			expectedResult: gopoet.Package{ImportPath: "foo", Name: "foo_bar"},
+		},
+		{
+			fd: mustBuildFile(builder.NewFile("foo/test.proto").
+				SetPackageName("foo.bar").
+				SetOptions(&descriptor.FileOptions{GoPackage: proto.String("foo.com/bar")})),
+			expectedResult: gopoet.Package{ImportPath: "foo.com/bar", Name: "bar"},
+		},
+		{
+			fd: mustBuildFile(builder.NewFile("foo/test.proto").
+				SetPackageName("foo.bar").
+				SetOptions(&descriptor.FileOptions{GoPackage: proto.String("foo.com/bar;bar_pkg")})),
+			expectedResult: gopoet.Package{ImportPath: "foo.com/bar", Name: "bar_pkg"},
+		},
+		{
+			fd: mustBuildFile(builder.NewFile("foo/test.proto").
+				SetPackageName("foo.bar").
+				SetOptions(&descriptor.FileOptions{GoPackage: proto.String("foo.com/bar;bar_pkg")})),
+			importMap:      map[string]string{"foo/blah.proto": "foo.io/baz"}, // not a match
+			expectedResult: gopoet.Package{ImportPath: "foo.com/bar", Name: "bar_pkg"},
+		},
+		{
+			fd: mustBuildFile(builder.NewFile("foo/test.proto").
+				SetPackageName("foo.bar").
+				SetOptions(&descriptor.FileOptions{GoPackage: proto.String("foo.com/bar;bar_pkg")})),
+			importMap:      map[string]string{"foo/test.proto": "foo.io/baz"},
+			expectedResult: gopoet.Package{ImportPath: "foo.io/baz", Name: "baz"},
+		},
+		{
+			fd: mustBuildFile(builder.NewFile("foo/test.proto").
+				SetPackageName("foo.bar").
+				SetOptions(&descriptor.FileOptions{GoPackage: proto.String("foo.com/bar;bar_pkg")})),
+			importMap:      map[string]string{"foo/test.proto": "foo.io/baz"},
+			override:       "foo.net/bar/baz;bar_baz",
+			expectedResult: gopoet.Package{ImportPath: "foo.net/bar/baz", Name: "bar_baz"},
+		},
+	}
+
+	for i, testCase := range testCases {
+		n := GoNames{ImportMap: testCase.importMap}
+		var pkg gopoet.Package
+		if testCase.override != "" {
+			pkg = n.GoPackageForFileWithOverride(testCase.fd, testCase.override)
+		} else {
+			pkg = n.GoPackageForFile(testCase.fd)
+		}
+		if pkg != testCase.expectedResult {
+			t.Errorf("case %d: expected package %s %q, got %s %q", i, testCase.expectedResult.Name, testCase.expectedResult.ImportPath, pkg.Name, pkg.ImportPath)
+		}
+
+		if testCase.override != "" {
+			// result is cached, so make sure we can get same package on
+			// subsequent call, even if we don't supply the same override
+			pkg2 := n.GoPackageForFile(testCase.fd)
+			if pkg2 != pkg {
+				t.Errorf("case %d: package with override %q not correctly cached", i, testCase.override)
+			}
+		}
+	}
+}
+
+func TestOutputFilenameFor(t *testing.T) {
+	// TODO
+}
+
+func TestGoNameOfField(t *testing.T) {
+	// TODO
+}
+
+func TestGoNameOfMethod(t *testing.T) {
+	// TODO
+}
+
+func TestGoNameOfServiceDesc(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForEnum(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForMessage(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForOneof(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForOneofChoice(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForServiceClient(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForStreamClient(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForStreamClientImpl(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForServiceServer(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForStreamServer(t *testing.T) {
+	// TODO
+}
+
+func TestGoTypeForStreamServerImpl(t *testing.T) {
+	// TODO
+}