Merge 34ebfb3 into 8772bad

codegen/client.go

@@ -258,7 +258,7 @@ func newGRPCClientSpec(
 	h *PackageHelper,
 ) (*ClientSpec, error) {
 	protoFile := filepath.Join(h.ThriftIDLPath(), config.IDLFile)
-	protoSpec, err := NewProtoModuleSpec(protoFile, false)
+	protoSpec, err := NewProtoModuleSpec(protoFile, false, h)
 	if err != nil {
 		return nil, errors.Wrapf(
 			err, "could not build proto spec for proto file %s: ", protoFile,

codegen/module_system.go

@@ -871,23 +871,16 @@ func (g *GRPCClientGenerator) Generate(
 		return nil, err
 	}
 
-	parts := strings.Split(clientSpec.ThriftFile, "/")
-	genDir := strings.Join(parts[len(parts)-3:len(parts)-1], "/")
-	genPkg := filepath.Join(
-		g.packageHelper.GenCodePackage(),
-		genDir,
-	)
 	// @rpatali: Update all struct to use more general field IDLFile instead of thriftFile.
 	clientMeta := &ClientMeta{
 		ProtoServices:    clientSpec.ModuleSpec.ProtoServices,
 		Instance:         instance,
 		ExportName:       clientSpec.ExportName,
 		ExportType:       clientSpec.ExportType,
 		Services:         nil,
-		IncludedPackages: nil,
+		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
 		ClientID:         clientSpec.ClientID,
 		ExposedMethods:   reversedMethods,
-		GenPkg:           genPkg,
 	}
 
 	client, err := g.templates.ExecTemplate(
@@ -1533,7 +1526,6 @@ type ClientMeta struct {
 	SidecarRouter    string
 	Fixture          *Fixture
 	DeputyReqHeader  string
-	GenPkg           string
 }
 
 func findMethod(

codegen/package.go

@@ -29,6 +29,11 @@ import (
 	"go.uber.org/thriftrw/compile"
 )
 
+const (
+	_thriftSuffix = ".thrift"
+	_protoSuffix  = ".proto"
+)
+
 // PackageHelper manages the mapping from thrift file to generated type code and service code.
 type PackageHelper struct {
 	// The project package name
@@ -236,20 +241,29 @@ func (p PackageHelper) GenCodePackage() string {
 
 // TypeImportPath returns the Go import path for types defined in a thrift file.
 func (p PackageHelper) TypeImportPath(thrift string) (string, error) {
-	if !strings.HasSuffix(thrift, ".thrift") {
-		return "", errors.Errorf("file %s is not .thrift", thrift)
+	if !strings.HasSuffix(thrift, _thriftSuffix) && !strings.HasSuffix(thrift, _protoSuffix) {
+		return "", errors.Errorf("file %s is not %s or %s", thrift, _thriftSuffix, _protoSuffix)
+	}
+	var suffix string
+	// for a filepath: a/b/c.(thrift|proto)
+	// - thrift generates code in path: a/b/c/c.go
+	// - proto generates code in path: a/b/c.go
+	if strings.HasSuffix(thrift, _thriftSuffix) {
+		suffix = strings.TrimSuffix(thrift, _thriftSuffix)
+	} else {
+		suffix = filepath.Dir(thrift)
 	}
 
 	idx := strings.Index(thrift, p.thriftRootDir)
 	if idx == -1 {
 		return "", errors.Errorf(
-			"file %s is not in thrift dir (%s)",
+			"file %s is not in IDL dir (%s)",
 			thrift, p.thriftRootDir,
 		)
 	}
 	return path.Join(
 		p.genCodePackage,
-		thrift[idx+len(p.thriftRootDir):len(thrift)-7],
+		thrift[idx+len(p.thriftRootDir):len(suffix)],
 	), nil
 }
 
@@ -271,32 +285,38 @@ func (p PackageHelper) CodeGenTargetPath() string {
 
 // TypePackageName returns the package name that defines the type.
 func (p PackageHelper) TypePackageName(thrift string) (string, error) {
-	if !strings.HasSuffix(thrift, ".thrift") {
-		return "", errors.Errorf("file %s is not .thrift", thrift)
+	if !strings.HasSuffix(thrift, _thriftSuffix) && !strings.HasSuffix(thrift, _protoSuffix) {
+		return "", errors.Errorf("file %s is not %s or %s", thrift, _thriftSuffix, _protoSuffix)
 	}
 	idx := strings.Index(thrift, p.thriftRootDir)
 	if idx == -1 {
 		return "", errors.Errorf(
-			"file %s is not in thrift dir (%s)",
+			"file %s is not in IDL dir (%s)",
 			thrift, p.thriftRootDir,
 		)
 	}
+	var suffix string
+	if strings.HasSuffix(thrift, _thriftSuffix) {
+		suffix = strings.TrimSuffix(thrift, _thriftSuffix)
+	} else {
+		suffix = strings.TrimSuffix(thrift, _protoSuffix)
+	}
 
 	// Strip the leading / and strip the .thrift on the end.
-	thriftSegment := thrift[idx+len(p.thriftRootDir)+1 : len(thrift)-7]
+	thriftSegment := thrift[idx+len(p.thriftRootDir)+1 : len(suffix)]
 
 	thriftPackageName := strings.Replace(thriftSegment, "/", "_", -1)
 	return CamelCase(thriftPackageName), nil
 }
 
 func (p PackageHelper) getRelativeFileName(thrift string) (string, error) {
-	if !strings.HasSuffix(thrift, ".thrift") {
-		return "", errors.Errorf("file %s is not .thrift", thrift)
+	if !strings.HasSuffix(thrift, _thriftSuffix) && !strings.HasSuffix(thrift, _protoSuffix) {
+		return "", errors.Errorf("file %s is not %s or %s", thrift, _thriftSuffix, _protoSuffix)
 	}
 	idx := strings.Index(thrift, p.thriftRootDir)
 	if idx == -1 {
 		return "", errors.Errorf(
-			"file %s is not in thrift dir (%s)",
+			"file %s is not in IDL dir (%s)",
 			thrift, p.thriftRootDir,
 		)
 	}

codegen/proto.go

@@ -24,6 +24,14 @@ import (
 	"github.com/emicklei/proto"
 )
 
+// ProtoModule is an internal representation of a parsed Proto file.
+type ProtoModule struct {
+	PackageName string
+	FilePath    string
+	Imports     []string
+	Services    []*ProtoService
+}
+
 // ProtoService is an internal representation of Proto service and methods in that service.
 type ProtoService struct {
 	Name string
@@ -43,24 +51,29 @@ type ProtoMessage struct {
 }
 
 type visitor struct {
-	protoServices []*ProtoService
+	Module *ProtoModule
 }
 
 func newVisitor() *visitor {
 	return &visitor{
-		protoServices: make([]*ProtoService, 0),
+		Module: new(ProtoModule),
 	}
 }
 
-func (v *visitor) Visit(proto *proto.Proto) []*ProtoService {
+func (v *visitor) Visit(proto *proto.Proto) *ProtoModule {
+	v.Module = &ProtoModule{
+		FilePath: proto.Filename,
+		Services: make([]*ProtoService, 0),
+		Imports:  make([]string, 0),
+	}
 	for _, e := range proto.Elements {
 		e.Accept(v)
 	}
-	return v.protoServices
+	return v.Module
 }
 
 func (v *visitor) VisitService(e *proto.Service) {
-	v.protoServices = append(v.protoServices, &ProtoService{
+	v.Module.Services = append(v.Module.Services, &ProtoService{
 		Name: e.Name,
 		RPC:  make([]*ProtoRPC, 0),
 	})
@@ -70,23 +83,29 @@ func (v *visitor) VisitService(e *proto.Service) {
 }
 
 func (v *visitor) VisitRPC(r *proto.RPC) {
-	s := v.protoServices[len(v.protoServices)-1]
+	s := v.Module.Services[len(v.Module.Services)-1]
 	s.RPC = append(s.RPC, &ProtoRPC{
 		Name:     r.Name,
 		Request:  &ProtoMessage{Name: r.RequestType},
 		Response: &ProtoMessage{Name: r.ReturnsType},
 	})
 }
 
+func (v *visitor) VisitPackage(e *proto.Package) {
+	v.Module.PackageName = e.Name
+}
+
+func (v *visitor) VisitImport(e *proto.Import) {
+	v.Module.Imports = append(v.Module.Imports, e.Filename)
+}
+
 // From the current use case, the following visits are no-op
 // since we only require the service, rpc methods and the request/response
 // types of those methods.
 
 func (v *visitor) VisitMessage(e *proto.Message)         {}
 func (v *visitor) VisitSyntax(e *proto.Syntax)           {}
-func (v *visitor) VisitPackage(e *proto.Package)         {}
 func (v *visitor) VisitOption(e *proto.Option)           {}
-func (v *visitor) VisitImport(e *proto.Import)           {}
 func (v *visitor) VisitNormalField(e *proto.NormalField) {}
 func (v *visitor) VisitEnumField(e *proto.EnumField)     {}
 func (v *visitor) VisitEnum(e *proto.Enum)               {}

codegen/proto_test.go

@@ -86,71 +86,85 @@ const (
 )
 
 var (
-	singleServiceSpecList = []*ProtoService{{
-		Name: "EchoService",
-		RPC: []*ProtoRPC{
-			{
-				Name: "EchoMethod",
-				Request: &ProtoMessage{
-					Name: "Request",
-				},
-				Response: &ProtoMessage{
-					Name: "Response",
-				},
-			},
-		},
-	}}
-	multiServiceSpecList = []*ProtoService{{
-		Name: "EchoService",
-		RPC: []*ProtoRPC{
-			{
-				Name: "EchoMethod1",
-				Request: &ProtoMessage{
-					Name: "Request1",
-				},
-				Response: &ProtoMessage{
-					Name: "Response1",
-				},
-			},
-			{
-				Name: "EchoMethod2",
-				Request: &ProtoMessage{
-					Name: "Request2",
-				},
-				Response: &ProtoMessage{
-					Name: "Response2",
+	singleServiceSpecList = &ProtoModule{
+		PackageName: "echo",
+		Services: []*ProtoService{{
+			Name: "EchoService",
+			RPC: []*ProtoRPC{
+				{
+					Name: "EchoMethod",
+					Request: &ProtoMessage{
+						Name: "Request",
+					},
+					Response: &ProtoMessage{
+						Name: "Response",
+					},
 				},
 			},
-		},
-	}}
-	mixedServiceSpecList = []*ProtoService{{
-		Name: "EchoService",
-		RPC: []*ProtoRPC{
-			{
-				Name: "EchoMethod1",
-				Request: &ProtoMessage{
-					Name: "Request1",
+		}},
+	}
+	multiServiceSpecList = &ProtoModule{
+		PackageName: "echo",
+		Services: []*ProtoService{{
+			Name: "EchoService",
+			RPC: []*ProtoRPC{
+				{
+					Name: "EchoMethod1",
+					Request: &ProtoMessage{
+						Name: "Request1",
+					},
+					Response: &ProtoMessage{
+						Name: "Response1",
+					},
 				},
-				Response: &ProtoMessage{
-					Name: "Response1",
+				{
+					Name: "EchoMethod2",
+					Request: &ProtoMessage{
+						Name: "Request2",
+					},
+					Response: &ProtoMessage{
+						Name: "Response2",
+					},
 				},
 			},
-			{
-				Name: "EchoMethod2",
-				Request: &ProtoMessage{
-					Name: "Request1",
+		}},
+	}
+	mixedServiceSpecList = &ProtoModule{
+		PackageName: "echo",
+		Services: []*ProtoService{{
+			Name: "EchoService",
+			RPC: []*ProtoRPC{
+				{
+					Name: "EchoMethod1",
+					Request: &ProtoMessage{
+						Name: "Request1",
+					},
+					Response: &ProtoMessage{
+						Name: "Response1",
+					},
 				},
-				Response: &ProtoMessage{
-					Name: "Response2",
+				{
+					Name: "EchoMethod2",
+					Request: &ProtoMessage{
+						Name: "Request1",
+					},
+					Response: &ProtoMessage{
+						Name: "Response2",
+					},
 				},
 			},
-		},
-	}}
-	noServiceSpecList    = make([]*ProtoService, 0)
-	emptyServiceSpecList = []*ProtoService{{
-		Name: "EchoService",
-		RPC:  make([]*ProtoRPC, 0),
-	}}
+		}},
+	}
+	noServiceSpecList = &ProtoModule{
+		PackageName: "echo",
+	}
+	emptyServiceSpecList = &ProtoModule{
+		PackageName: "echo",
+		Services: []*ProtoService{{
+			Name: "EchoService",
+			RPC:  make([]*ProtoRPC, 0),
+		}},
+	}
 )
 
 func TestRunner(t *testing.T) {
@@ -161,12 +175,13 @@ func TestRunner(t *testing.T) {
 	assertElementMatch(t, emptyServiceSpec, emptyServiceSpecList)
 }
 
-func assertElementMatch(t *testing.T, specRaw string, specParsed []*ProtoService) {
+func assertElementMatch(t *testing.T, specRaw string, specParsed *ProtoModule) {
 	r := strings.NewReader(specRaw)
 	parser := proto.NewParser(r)
 	protoSpec, err := parser.Parse()
 	assert.NoErrorf(t, err, "proto spec parsing failed")
 
 	v := newVisitor().Visit(protoSpec)
-	assert.ElementsMatch(t, v, specParsed)
+	assert.Equal(t, specParsed.PackageName, v.PackageName)
+	assert.ElementsMatch(t, specParsed.Services, v.Services)
 }