Repository: googleapis/gapic-generator-go
Pull request #140 (2019-06-05 21:53:52)

SampleGen: iterate over a map field in response handling
Porting https://github.com/googleapis/gapic-generator/pull/2593 over to Go.

hzyi-google: PTAL

hzyi-google: PTAL

Patch:
diff --git a/cmd/gen-go-sample/gapic.go b/cmd/gen-go-sample/gapic.go
index 5a37fd3..bf9705f 100644
--- a/cmd/gen-go-sample/gapic.go
+++ b/cmd/gen-go-sample/gapic.go
@@ -69,7 +69,10 @@ type OutputSpec struct {
 
 type LoopSpec struct {
 	Collection string
+	Map        string
 	Variable   string
+	Key        string
+	Value      string
 	Body       []OutputSpec
 }
 
diff --git a/cmd/gen-go-sample/inittree.go b/cmd/gen-go-sample/inittree.go
index 376bcea..892eb4e 100644
--- a/cmd/gen-go-sample/inittree.go
+++ b/cmd/gen-go-sample/inittree.go
@@ -43,6 +43,10 @@ type initType struct {
 
 	repeated bool
 
+	// keyType and valueType are set if the type is a protobuf map
+	keyType   *initType
+	valueType *initType
+
 	// valFmt, if not nil, post-processes values to be included into init struct.
 	// NOTE(pongad): This func signature might seem too general. I think it is just general enough
 	// to deal with enums and bytes. Time will tell.
@@ -106,6 +110,35 @@ func (t *initTree) get(k string, info pbinfo.Info) (*initTree, error) {
 		} else {
 			// type is a message
 			v.typ.desc = typ
+
+			if d, ok := typ.(*descriptor.DescriptorProto); ok {
+				if d.GetOptions().GetMapEntry() == true {
+					// type is a map entry
+					for _, f2 := range d.Field {
+						switch f2.GetName() {
+						case "key":
+							// key is always a primitive type
+							v.typ.keyType = &initType{prim: f2.GetType()}
+						case "value":
+							t2 := f2.GetType()
+
+							// value is a primitive type
+							v.typ.valueType = &initType{prim: t2}
+
+							if t2 == descriptor.FieldDescriptorProto_TYPE_MESSAGE || t2 == descriptor.FieldDescriptorProto_TYPE_ENUM {
+								// value is a protobuf message or enum
+								tn2 := f2.GetTypeName()
+								typ2 := info.Type[tn2]
+								if typ2 == nil {
+									return nil, errors.E(nil, "cannot find descriptor of %q", tn2)
+								}
+
+								v.typ.valueType = &initType{desc: typ2}
+							}
+						}
+					}
+				}
+			}
 		}
 		v.typ.repeated = f.GetLabel() == descriptor.FieldDescriptorProto_LABEL_REPEATED
 
diff --git a/cmd/gen-go-sample/out.go b/cmd/gen-go-sample/out.go
index 6bd2bad..cbd53d5 100644
--- a/cmd/gen-go-sample/out.go
+++ b/cmd/gen-go-sample/out.go
@@ -86,7 +86,14 @@ func writeOutputSpec(out OutputSpec, st *symTab, gen *generator) error {
 	}
 	if l := out.Loop; l != nil {
 		used++
-		err = writeLoop(l, st, gen)
+		err = errors.E(nil, "")
+		if l.Collection != "" && l.Map != "" {
+			err = errors.E(nil, "only one of collection and map should be set")
+		} else if l.Collection != "" {
+			err = writeLoop(l, st, gen)
+		} else if l.Map != "" {
+			err = writeMap(l, st, gen)
+		}
 	}
 
 	if used == 0 {
@@ -133,7 +140,7 @@ func writePrint(pFmt string, pArgs []string, st *symTab, gen *generator) error {
 		sc, report := initScanner(arg)
 		path, _, err := writePath(sc, st, gen.descInfo)
 		if err != nil {
-			return report(nil)
+			return report(err)
 		}
 		sb.WriteString(path)
 	}
@@ -145,6 +152,10 @@ func writePrint(pFmt string, pArgs []string, st *symTab, gen *generator) error {
 }
 
 func writeLoop(l *LoopSpec, st *symTab, gen *generator) error {
+	if l.Variable == "" {
+		return errors.E(nil, "variable not specified for looping over arrays")
+	}
+
 	p := gen.pt.Printf
 
 	sc, report := initScanner(l.Collection)
@@ -168,6 +179,48 @@ func writeLoop(l *LoopSpec, st *symTab, gen *generator) error {
 	return nil
 }
 
+func writeMap(l *LoopSpec, st *symTab, gen *generator) error {
+	if l.Key == "" && l.Value == "" {
+		return errors.E(nil, "at least one of key and value should be specified for looping over maps")
+	}
+
+	p := gen.pt.Printf
+
+	sc, report := initScanner(l.Map)
+	path, typ, err := writePath(sc, st, gen.descInfo)
+	if err = report(err); err != nil {
+		return err
+	}
+
+	if typ.keyType == nil || typ.valueType == nil {
+		return errors.E(nil, "%s is not a map field", l.Map)
+	}
+	keyType := *typ.keyType
+	valueType := *typ.valueType
+
+	stInner := newSymTab(st)
+
+	if l.Value == "" {
+		p("for %s := range %s {", snakeToCamel(l.Key), path)
+		stInner.put(l.Key, keyType)
+	} else if l.Key == "" {
+		p("for _, %s := range %s {", snakeToCamel(l.Value), path)
+		stInner.put(l.Value, valueType)
+	} else {
+		p("for %s, %s := range %s {", snakeToCamel(l.Key), snakeToCamel(l.Value), path)
+		stInner.put(l.Key, keyType)
+		stInner.put(l.Value, valueType)
+	}
+
+	for _, b := range l.Body {
+		if err := writeOutputSpec(b, stInner, gen); err != nil {
+			return err
+		}
+	}
+	p("}")
+	return nil
+}
+
 func writePath(sc *scanner.Scanner, st *symTab, info pbinfo.Info) (string, initType, error) {
 	if sc.Scan() != scanner.Ident {
 		return "", initType{}, errors.E(nil, "expecting ident, got %s", sc.TokenText())


Test patch:
diff --git a/cmd/gen-go-sample/main_test.go b/cmd/gen-go-sample/main_test.go
index c344153..1f2ac31 100644
--- a/cmd/gen-go-sample/main_test.go
+++ b/cmd/gen-go-sample/main_test.go
@@ -20,6 +20,7 @@ import (
 
 	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/protoc-gen-go/descriptor"
+	"github.com/googleapis/gapic-generator-go/internal/errors"
 	"github.com/googleapis/gapic-generator-go/internal/pbinfo"
 	"github.com/googleapis/gapic-generator-go/internal/txtdiff"
 )
@@ -186,16 +187,82 @@ func TestEmpty(t *testing.T) {
 	compare(t, g, filepath.Join("testdata", "sample_empty.want"))
 }
 
+func TestMapOut(t *testing.T) {
+	t.Parallel()
+
+	g := initTestGenerator()
+	vs := SampleValueSet{
+		ID: "my_value_set",
+		OnSuccess: []OutputSpec{
+			{
+				Loop: &LoopSpec{
+					Key: "just_key",
+					Map: "$resp.mappy_map",
+					Body: []OutputSpec{
+						{Print: []string{"key: %s", "just_key"}},
+					},
+				},
+			},
+			{
+				Loop: &LoopSpec{
+					Key:   "k",
+					Value: "v",
+					Map:   "$resp.mappy_map",
+					Body: []OutputSpec{
+						{Print: []string{"key: %s, value: %s", "k", "v"}},
+					},
+				},
+			},
+			{
+				Loop: &LoopSpec{
+					Value: "only_value",
+					Map:   "$resp.mappy_map",
+					Body: []OutputSpec{
+						{Print: []string{"value: %s", "only_value"}},
+					},
+				},
+			},
+		},
+	}
+	if err := g.genSample("foo.FooService", GAPICMethod{Name: "UnaryMethod"}, "awesome_region", vs); err != nil {
+		t.Fatal(err)
+	}
+	compare(t, g, filepath.Join("testdata", "sample_map_out.want"))
+}
+
 func initTestGenerator() *generator {
-	labelp := func(l descriptor.FieldDescriptorProto_Label) *descriptor.FieldDescriptorProto_Label {
-		return &l
+	eType := &descriptor.EnumDescriptorProto{
+		Name: proto.String("EType"),
+		Value: []*descriptor.EnumValueDescriptorProto{
+			{Name: proto.String("FOO")},
+		},
 	}
 
+	aType := &descriptor.DescriptorProto{
+		Name: proto.String("AType"),
+		Field: []*descriptor.FieldDescriptorProto{
+			{Name: proto.String("x"), Type: typep(descriptor.FieldDescriptorProto_TYPE_INT64)},
+			{Name: proto.String("y"), Type: typep(descriptor.FieldDescriptorProto_TYPE_FLOAT)},
+		},
+		EnumType: []*descriptor.EnumDescriptorProto{eType},
+	}
+
+	mapType, mapField :=
+		createMapTypeAndField(
+			"mappy_map",
+			".foo.InputType",
+			typep(descriptor.FieldDescriptorProto_TYPE_STRING),
+			typep(descriptor.FieldDescriptorProto_TYPE_MESSAGE),
+			".foo.AType")
+
 	inType := &descriptor.DescriptorProto{
 		Name: proto.String("InputType"),
 		OneofDecl: []*descriptor.OneofDescriptorProto{
 			{Name: proto.String("Group")},
 		},
+		NestedType: []*descriptor.DescriptorProto{
+			mapType,
+		},
 		Field: []*descriptor.FieldDescriptorProto{
 			{Name: proto.String("a"), TypeName: proto.String(".foo.AType")},
 			{Name: proto.String("a_array"), TypeName: proto.String(".foo.AType"), Label: labelp(descriptor.FieldDescriptorProto_LABEL_REPEATED)},
@@ -207,22 +274,9 @@ func initTestGenerator() *generator {
 			{Name: proto.String("r"), Type: typep(descriptor.FieldDescriptorProto_TYPE_STRING), Label: labelp(descriptor.FieldDescriptorProto_LABEL_REPEATED)},
 			{Name: proto.String("resource_field"), Type: typep(descriptor.FieldDescriptorProto_TYPE_STRING)},
 			{Name: proto.String("bytes"), Type: typep(descriptor.FieldDescriptorProto_TYPE_BYTES)},
+			mapField,
 		},
 	}
-	eType := &descriptor.EnumDescriptorProto{
-		Name: proto.String("EType"),
-		Value: []*descriptor.EnumValueDescriptorProto{
-			{Name: proto.String("FOO")},
-		},
-	}
-	aType := &descriptor.DescriptorProto{
-		Name: proto.String("AType"),
-		Field: []*descriptor.FieldDescriptorProto{
-			{Name: proto.String("x"), Type: typep(descriptor.FieldDescriptorProto_TYPE_INT64)},
-			{Name: proto.String("y"), Type: typep(descriptor.FieldDescriptorProto_TYPE_FLOAT)},
-		},
-		EnumType: []*descriptor.EnumDescriptorProto{eType},
-	}
 
 	pageInType := &descriptor.DescriptorProto{
 		Name: proto.String("PageInType"),
@@ -299,6 +353,61 @@ func initTestGenerator() *generator {
 	}
 }
 
+func labelp(l descriptor.FieldDescriptorProto_Label) *descriptor.FieldDescriptorProto_Label {
+	return &l
+}
+
+// createMapTypeAndField creates the generated MapEntry protobuf message and the actual map field.
+// If valueType is enum or message, vTypeName must not be empty.
+func createMapTypeAndField(
+	fieldName string,
+	parentTyp string,
+	keyType *descriptor.FieldDescriptorProto_Type,
+	valueType *descriptor.FieldDescriptorProto_Type,
+	vTypeName string,
+) (*descriptor.DescriptorProto, *descriptor.FieldDescriptorProto) {
+
+	var vf *descriptor.FieldDescriptorProto
+	if vTypeName == "" {
+		if *valueType == descriptor.FieldDescriptorProto_TYPE_ENUM {
+			panic(errors.E(nil, "expecting non-empty enum type name"))
+		}
+		if *valueType == descriptor.FieldDescriptorProto_TYPE_MESSAGE {
+			panic(errors.E(nil, "expecting non-empty message type name"))
+		}
+		vf = &descriptor.FieldDescriptorProto{
+			Name: proto.String("value"),
+			Type: valueType,
+		}
+	} else {
+		vf = &descriptor.FieldDescriptorProto{
+			Name:     proto.String("value"),
+			Type:     valueType,
+			TypeName: proto.String(vTypeName),
+		}
+	}
+
+	mapEntry := &descriptor.DescriptorProto{
+		Name: proto.String("MapFieldEntry"),
+		Field: []*descriptor.FieldDescriptorProto{
+			{Name: proto.String("key"), Type: keyType},
+			vf,
+		},
+		Options: &descriptor.MessageOptions{
+			MapEntry: proto.Bool(true),
+		},
+	}
+
+	mapField := &descriptor.FieldDescriptorProto{
+		Name:     proto.String(fieldName),
+		Type:     typep(descriptor.FieldDescriptorProto_TYPE_MESSAGE),
+		TypeName: proto.String(parentTyp + ".MapFieldEntry"),
+		Label:    labelp(descriptor.FieldDescriptorProto_LABEL_REPEATED),
+	}
+
+	return mapEntry, mapField
+}
+
 func compare(t *testing.T, g *generator, goldenPath string) {
 	t.Helper()
 
diff --git a/cmd/gen-go-sample/testdata/sample_map_out.want b/cmd/gen-go-sample/testdata/sample_map_out.want
new file mode 100644
index 0000000..92874bd
--- /dev/null
+++ b/cmd/gen-go-sample/testdata/sample_map_out.want
@@ -0,0 +1,63 @@
+// Copyright 2018 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Code generated by protoc-gen-go_gapic. DO NOT EDIT.
+
+package main
+import(
+ "context"
+ "flag"
+ "fmt"
+ "log"
+
+foo "path.to/client/foo"
+foopb "path.to/pb/foo"
+)
+// [START awesome_region]
+
+func sampleUnaryMethod() error {
+	ctx := context.Background()
+	c, err := foo.NewClient(ctx)
+	if err != nil {
+		return err
+	}
+
+	req := &foopb.InputType{
+	}
+	resp, err := c.UnaryMethod(ctx, req)
+	if err != nil {
+		return err
+	}
+
+	for justKey := range resp.GetMappyMap() {
+		fmt.Printf("key: %v\n", justKey)
+	}
+	for k, v := range resp.GetMappyMap() {
+		fmt.Printf("key: %v, value: %v\n", k, v)
+	}
+	for _, onlyValue := range resp.GetMappyMap() {
+		fmt.Printf("value: %v\n", onlyValue)
+	}
+	return nil
+}
+
+// [END awesome_region]
+
+func main() {
+	flag.Parse()
+	if err := sampleUnaryMethod(); err != nil {
+		log.Fatal(err)
+	}
+}
+