Repository: googleapis/gapic-generator-go
Pull request #14 (2018-09-25 00:44:39)

pbinfo: add enums to lookup map
This doesn't benefit the generated clients but will benefit
the generated samples.

The vast majority of our code doesn't care for the difference between
a message type and an enum type, only that we have a name to refer to
them.

pongad: Thank you for reviews @jba, PTAL

Patch:
diff --git a/internal/gengapic/example.go b/internal/gengapic/example.go
index 2400ff4..d5c7d11 100644
--- a/internal/gengapic/example.go
+++ b/internal/gengapic/example.go
@@ -70,7 +70,7 @@ func (g *generator) exampleMethod(pkgName, servName string, m *descriptor.Method
 
 	if !m.GetClientStreaming() && !m.GetServerStreaming() {
 		p("")
-		p("req := &%s.%s{", inSpec.Name, *inType.Name)
+		p("req := &%s.%s{", inSpec.Name, inType.GetName())
 		p("  // TODO: Fill request struct fields.")
 		p("}")
 	}
@@ -150,7 +150,7 @@ func (g *generator) examplePagingCall(m *descriptor.MethodDescriptorProto) {
 	g.imports[pbinfo.ImportSpec{Path: "google.golang.org/api/iterator"}] = true
 }
 
-func (g *generator) exampleBidiCall(m *descriptor.MethodDescriptorProto, inType *descriptor.DescriptorProto, inSpec pbinfo.ImportSpec) {
+func (g *generator) exampleBidiCall(m *descriptor.MethodDescriptorProto, inType pbinfo.ProtoType, inSpec pbinfo.ImportSpec) {
 	p := g.printf
 
 	p("stream, err := c.%s(ctx)", m.GetName())
diff --git a/internal/gengapic/gengapic.go b/internal/gengapic/gengapic.go
index f0320bb..9b09759 100644
--- a/internal/gengapic/gengapic.go
+++ b/internal/gengapic/gengapic.go
@@ -344,11 +344,11 @@ func (g *generator) unaryCall(servName string, m *descriptor.MethodDescriptorPro
 	p := g.printf
 
 	p("func (c *%sClient) %s(ctx context.Context, req *%s.%s, opts ...gax.CallOption) (*%s.%s, error) {",
-		servName, *m.Name, inSpec.Name, *inType.Name, outSpec.Name, *outType.Name)
+		servName, *m.Name, inSpec.Name, inType.GetName(), outSpec.Name, outType.GetName())
 
 	g.insertMetadata()
 	g.appendCallOpts(m)
-	p("var resp *%s.%s", outSpec.Name, *outType.Name)
+	p("var resp *%s.%s", outSpec.Name, outType.GetName())
 	p("err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {")
 	p("  var err error")
 	p("  resp, err = %s", grpcClientCall(servName, *m.Name))
diff --git a/internal/gengapic/lro.go b/internal/gengapic/lro.go
index a1eca6d..67008d8 100644
--- a/internal/gengapic/lro.go
+++ b/internal/gengapic/lro.go
@@ -43,11 +43,11 @@ func (g *generator) lroCall(servName string, m *descriptor.MethodDescriptorProto
 	p := g.printf
 
 	p("func (c *%sClient) %s(ctx context.Context, req *%s.%s, opts ...gax.CallOption) (*%s, error) {",
-		servName, *m.Name, inSpec.Name, *inType.Name, lroType)
+		servName, *m.Name, inSpec.Name, inType.GetName(), lroType)
 
 	g.insertMetadata()
 	g.appendCallOpts(m)
-	p("  var resp *%s.%s", outSpec.Name, *outType.Name)
+	p("  var resp *%s.%s", outSpec.Name, outType.GetName())
 	p("  err := gax.Invoke(ctx, func(ctx context.Context, settings gax.CallSettings) error {")
 	p("    var err error")
 	p("    resp, err = %s", grpcClientCall(servName, *m.Name))
diff --git a/internal/gengapic/paging.go b/internal/gengapic/paging.go
index f3e9e1a..be25836 100644
--- a/internal/gengapic/paging.go
+++ b/internal/gengapic/paging.go
@@ -65,7 +65,7 @@ func (g *generator) iterTypeOf(elemField *descriptor.FieldDescriptorProto) (iter
 		}
 
 		pt.elemTypeName = fmt.Sprintf("*%s.%s", imp.Name, eType.GetName())
-		pt.iterTypeName = *eType.Name + "Iterator"
+		pt.iterTypeName = eType.GetName() + "Iterator"
 
 		pt.elemImports = []pbinfo.ImportSpec{imp}
 
@@ -101,12 +101,21 @@ func (g *generator) pagingField(m *descriptor.MethodDescriptorProto) (*descripto
 	if inType == nil {
 		return nil, errors.E(nil, "cannot find message type %q, malformed descriptor?", m.GetInputType())
 	}
+	inMsg, ok := inType.(*descriptor.DescriptorProto)
+	if !ok {
+		return nil, errors.E(nil, "expected %q to be message type, found %T", m.GetInputType(), inType)
+	}
+
 	outType := g.descInfo.Type[m.GetOutputType()]
 	if outType == nil {
 		return nil, errors.E(nil, "cannot find message type %q, malformed descriptor?", m.GetOutputType())
 	}
+	outMsg, ok := outType.(*descriptor.DescriptorProto)
+	if !ok {
+		return nil, errors.E(nil, "expected %q to be message type, found %T", m.GetOutputType(), outType)
+	}
 
-	for _, f := range inType.Field {
+	for _, f := range inMsg.Field {
 		if f.GetName() == "page_size" && f.GetType() == descriptor.FieldDescriptorProto_TYPE_INT32 {
 			hasSize = true
 		}
@@ -114,7 +123,7 @@ func (g *generator) pagingField(m *descriptor.MethodDescriptorProto) (*descripto
 			hasToken = true
 		}
 	}
-	for _, f := range outType.Field {
+	for _, f := range outMsg.Field {
 		if f.GetName() == "next_page_token" && f.GetType() == descriptor.FieldDescriptorProto_TYPE_STRING {
 			hasNextToken = true
 		}
@@ -126,10 +135,10 @@ func (g *generator) pagingField(m *descriptor.MethodDescriptorProto) (*descripto
 		return nil, nil
 	}
 	if len(elemFields) == 0 {
-		return nil, fmt.Errorf("%s looks like paging method, but can't find repeated field in %s", *m.Name, *outType.Name)
+		return nil, fmt.Errorf("%s looks like paging method, but can't find repeated field in %s", *m.Name, outType.GetName())
 	}
 	if len(elemFields) > 1 {
-		return nil, fmt.Errorf("%s looks like paging method, but too many repeated fields in %s", *m.Name, *outType.Name)
+		return nil, fmt.Errorf("%s looks like paging method, but too many repeated fields in %s", *m.Name, outType.GetName())
 	}
 	return elemFields[0], nil
 }
@@ -149,15 +158,15 @@ func (g *generator) pagingCall(servName string, m *descriptor.MethodDescriptorPr
 
 	p := g.printf
 	p("func (c *%sClient) %s(ctx context.Context, req *%s.%s, opts ...gax.CallOption) *%s {",
-		servName, *m.Name, inSpec.Name, *inType.Name, pt.iterTypeName)
+		servName, *m.Name, inSpec.Name, inType.GetName(), pt.iterTypeName)
 
 	g.insertMetadata()
 	g.appendCallOpts(m)
 
 	p("it := &%s{}", pt.iterTypeName)
-	p("req = proto.Clone(req).(*%s.%s)", inSpec.Name, *inType.Name)
+	p("req = proto.Clone(req).(*%s.%s)", inSpec.Name, inType.GetName())
 	p("it.InternalFetch = func(pageSize int, pageToken string) ([]%s, string, error) {", pt.elemTypeName)
-	p("  var resp *%s.%s", outSpec.Name, *outType.Name)
+	p("  var resp *%s.%s", outSpec.Name, outType.GetName())
 	p("  req.PageToken = pageToken")
 	p("  if pageSize > math.MaxInt32 {")
 	p("    req.PageSize = math.MaxInt32")
diff --git a/internal/pbinfo/pbinfo.go b/internal/pbinfo/pbinfo.go
index 74cb0b9..299dd30 100644
--- a/internal/pbinfo/pbinfo.go
+++ b/internal/pbinfo/pbinfo.go
@@ -24,13 +24,23 @@ import (
 	"github.com/googleapis/gapic-generator-go/internal/errors"
 )
 
+// ProtoType represents a type in protobuf descriptors.
+// It is an interface implemented by DescriptorProto and EnumDescriptorProto.
+type ProtoType interface {
+	proto.Message
+	GetName() string
+}
+
+// Info provides lookup tables for various protobuf properties.
+// For example, we can look up a type by name without iterating the entire
+// descriptor.
 type Info struct {
 	// Maps services and messages to the file containing them,
 	// so we can figure out the import.
 	ParentFile map[proto.Message]*descriptor.FileDescriptorProto
 
 	// Maps type names to their messages.
-	Type map[string]*descriptor.DescriptorProto
+	Type map[string]ProtoType
 
 	// Maps service names to their descriptors.
 	Serv map[string]*descriptor.ServiceDescriptorProto
@@ -40,7 +50,7 @@ type Info struct {
 func Of(files []*descriptor.FileDescriptorProto) Info {
 	info := Info{
 		ParentFile: map[proto.Message]*descriptor.FileDescriptorProto{},
-		Type:       map[string]*descriptor.DescriptorProto{},
+		Type:       map[string]ProtoType{},
 		Serv:       map[string]*descriptor.ServiceDescriptorProto{},
 	}
 
@@ -56,8 +66,10 @@ func Of(files []*descriptor.FileDescriptorProto) Info {
 		// Type
 		for _, m := range f.MessageType {
 			// In descriptors, putting the dot in front means the name is fully-qualified.
-			fullyQualifiedName := fmt.Sprintf(".%s.%s", f.GetPackage(), m.GetName())
-			info.Type[fullyQualifiedName] = m
+			addMessage(info.Type, "."+f.GetPackage(), m)
+		}
+		for _, e := range f.EnumType {
+			info.Type["."+f.GetPackage()+"."+e.GetName()] = e
 		}
 
 		// Serv
@@ -70,6 +82,18 @@ func Of(files []*descriptor.FileDescriptorProto) Info {
 	return info
 }
 
+func addMessage(m map[string]ProtoType, prefix string, msg *descriptor.DescriptorProto) {
+	fullName := prefix + "." + msg.GetName()
+	m[fullName] = msg
+
+	for _, subMsg := range msg.NestedType {
+		addMessage(m, fullName, subMsg)
+	}
+	for _, subEnum := range msg.EnumType {
+		m[fullName+"."+subEnum.GetName()] = subEnum
+	}
+}
+
 type ImportSpec struct {
 	Name, Path string
 }


Test patch:
diff --git a/cmd/gen-go-sample/main_test.go b/cmd/gen-go-sample/main_test.go
index e67903b..9d7d325 100644
--- a/cmd/gen-go-sample/main_test.go
+++ b/cmd/gen-go-sample/main_test.go
@@ -75,7 +75,7 @@ func TestSample(t *testing.T) {
 				serv.Method[0]: file,
 				inType:         file,
 			},
-			Type: map[string]*descriptor.DescriptorProto{
+			Type: map[string]pbinfo.ProtoType{
 				"InputType": inType,
 			},
 		},
diff --git a/internal/gengapic/example_test.go b/internal/gengapic/example_test.go
index 7d3d50b..5dda6ea 100644
--- a/internal/gengapic/example_test.go
+++ b/internal/gengapic/example_test.go
@@ -138,7 +138,7 @@ func commonTypes(g *generator) {
 		Name: proto.String("Operation"),
 	}
 
-	g.descInfo.Type = map[string]*descriptor.DescriptorProto{
+	g.descInfo.Type = map[string]pbinfo.ProtoType{
 		emptyType: empty,
 		lroType:   lro,
 	}
diff --git a/internal/gengapic/paging_test.go b/internal/gengapic/paging_test.go
index 62a2906..de8914d 100644
--- a/internal/gengapic/paging_test.go
+++ b/internal/gengapic/paging_test.go
@@ -38,7 +38,7 @@ func TestPagingField(t *testing.T) {
 	}
 
 	g := &generator{}
-	g.descInfo.Type = map[string]*descriptor.DescriptorProto{
+	g.descInfo.Type = map[string]pbinfo.ProtoType{
 		"Foo": &descriptor.DescriptorProto{
 			Name: proto.String("Foo"),
 		},
@@ -152,8 +152,8 @@ func TestIterTypeOf(t *testing.T) {
 	}
 	g := &generator{
 		descInfo: pbinfo.Info{
-			Type: map[string]*descriptor.DescriptorProto{
-				*msgType.Name: msgType,
+			Type: map[string]pbinfo.ProtoType{
+				msgType.GetName(): msgType,
 			},
 			ParentFile: map[proto.Message]*descriptor.FileDescriptorProto{
 				msgType: &descriptor.FileDescriptorProto{
@@ -190,7 +190,7 @@ func TestIterTypeOf(t *testing.T) {
 		{
 			field: &descriptor.FieldDescriptorProto{
 				Type:     typep(descriptor.FieldDescriptorProto_TYPE_MESSAGE),
-				TypeName: msgType.Name,
+				TypeName: proto.String(msgType.GetName()),
 			},
 			want: iterType{
 				iterTypeName: "FooIterator",