Skip to content

Commit 5ab7c00

Browse files
committed
1 parent 6177ba0 commit 5ab7c00

File tree

2 files changed

+272
-0
lines changed

2 files changed

+272
-0
lines changed

internal/pbinfo/pbinfo.go

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
// Copyright 2018 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
// Package pbinfo provides convenience types for looking up protobuf elements.
16+
package pbinfo
17+
18+
import (
19+
"fmt"
20+
"strings"
21+
"unicode"
22+
23+
"github.com/golang/protobuf/protoc-gen-go/descriptor"
24+
"google.golang.org/protobuf/proto"
25+
)
26+
27+
// ProtoType represents a type in protobuf descriptors.
28+
// It is an interface implemented by DescriptorProto and EnumDescriptorProto.
29+
type ProtoType interface {
30+
proto.Message
31+
GetName() string
32+
}
33+
34+
// Info provides lookup tables for various protobuf properties.
35+
// For example, we can look up a type by name without iterating the entire
36+
// descriptor.
37+
type Info struct {
38+
// Maps services and messages to the file containing them,
39+
// so we can figure out the import.
40+
ParentFile map[proto.Message]*descriptor.FileDescriptorProto
41+
42+
// NOTE(pongad): ParentElement and sub-types are only used in samples.
43+
// They are added in the shared package because they share a lot of similarities
44+
// with things that are already here. Maybe revisit this in the future?
45+
46+
// Maps a protobuf element to the enclosing scope.
47+
// If enum E is defined in message M which is in file F,
48+
// ParentElement[E]=M, ParentElement[M]=nil, and ParentFile[M]=F
49+
ParentElement map[ProtoType]ProtoType
50+
51+
// Maps type names to their messages.
52+
Type map[string]ProtoType
53+
54+
// Maps service names to their descriptors.
55+
Serv map[string]*descriptor.ServiceDescriptorProto
56+
57+
// PkgOverrides is file-to-import mapping used to override the
58+
// go_package option in the given proto file.
59+
PkgOverrides map[string]string
60+
}
61+
62+
// Of creates Info from given protobuf files.
63+
func Of(files []*descriptor.FileDescriptorProto) Info {
64+
info := Info{
65+
ParentFile: map[proto.Message]*descriptor.FileDescriptorProto{},
66+
ParentElement: map[ProtoType]ProtoType{},
67+
Type: map[string]ProtoType{},
68+
Serv: map[string]*descriptor.ServiceDescriptorProto{},
69+
PkgOverrides: map[string]string{},
70+
}
71+
72+
for _, f := range files {
73+
// ParentFile
74+
for _, m := range f.GetMessageType() {
75+
info.ParentFile[m] = f
76+
}
77+
for _, e := range f.GetEnumType() {
78+
info.ParentFile[e] = f
79+
}
80+
for _, s := range f.GetService() {
81+
info.ParentFile[s] = f
82+
for _, m := range s.GetMethod() {
83+
info.ParentFile[m] = f
84+
info.ParentElement[m] = s
85+
}
86+
}
87+
88+
// Type
89+
for _, m := range f.GetMessageType() {
90+
// In descriptors, putting the dot in front means the name is fully-qualified.
91+
addMessage(info.Type, info.ParentElement, "."+f.GetPackage(), m, nil)
92+
}
93+
for _, e := range f.GetEnumType() {
94+
info.Type["."+f.GetPackage()+"."+e.GetName()] = e
95+
}
96+
97+
// Serv
98+
for _, s := range f.GetService() {
99+
fullyQualifiedName := fmt.Sprintf(".%s.%s", f.GetPackage(), s.GetName())
100+
info.Serv[fullyQualifiedName] = s
101+
}
102+
}
103+
104+
return info
105+
}
106+
107+
func addMessage(typMap map[string]ProtoType, parentMap map[ProtoType]ProtoType, prefix string, msg, parentMsg *descriptor.DescriptorProto) {
108+
fullName := prefix + "." + msg.GetName()
109+
typMap[fullName] = msg
110+
if parentMsg != nil {
111+
parentMap[msg] = parentMsg
112+
}
113+
114+
for _, subMsg := range msg.GetNestedType() {
115+
addMessage(typMap, parentMap, fullName, subMsg, msg)
116+
}
117+
118+
for _, subEnum := range msg.GetEnumType() {
119+
typMap[fullName+"."+subEnum.GetName()] = subEnum
120+
parentMap[subEnum] = msg
121+
}
122+
123+
for _, field := range msg.GetField() {
124+
parentMap[field] = msg
125+
}
126+
}
127+
128+
// ImportSpec represents a Go module import path and an optional alias.
129+
type ImportSpec struct {
130+
Name, Path string
131+
}
132+
133+
// NameSpec reports the name and ImportSpec of e.
134+
//
135+
// The reported name is the same with how protoc-gen-go refers to e.
136+
// E.g. if type B is nested under A, then the name of type B is "A_B".
137+
func (in *Info) NameSpec(e ProtoType) (string, ImportSpec, error) {
138+
appendpb := func(n string) string {
139+
if !strings.HasSuffix(n, "pb") {
140+
n += "pb"
141+
}
142+
return n
143+
}
144+
145+
topLvl := e
146+
var nameParts []string
147+
for e2 := e; e2 != nil; e2 = in.ParentElement[e2] {
148+
topLvl = e2
149+
nameParts = append(nameParts, e2.GetName())
150+
}
151+
for i, l := 0, len(nameParts); i < l/2; i++ {
152+
nameParts[i], nameParts[l-i-1] = nameParts[l-i-1], nameParts[i]
153+
}
154+
name := strings.Join(nameParts, "_")
155+
156+
var eTxt interface{} = e
157+
if et, ok := eTxt.(interface{ GetName() string }); ok {
158+
eTxt = et.GetName()
159+
}
160+
161+
fdesc := in.ParentFile[topLvl]
162+
if fdesc == nil {
163+
return "", ImportSpec{}, fmt.Errorf("can't determine import path for %v; can't find parent file", eTxt)
164+
}
165+
166+
pkg := fdesc.GetOptions().GetGoPackage()
167+
if pkgOverride, ok := in.PkgOverrides[fdesc.GetName()]; ok {
168+
pkg = pkgOverride
169+
}
170+
if pkg == "" {
171+
return "", ImportSpec{}, fmt.Errorf("can't determine import path for %v, file %q missing `option go_package`", eTxt, fdesc.GetName())
172+
}
173+
174+
if p := strings.IndexByte(pkg, ';'); p >= 0 {
175+
return name, ImportSpec{Path: pkg[:p], Name: appendpb(pkg[p+1:])}, nil
176+
}
177+
178+
for {
179+
p := strings.LastIndexByte(pkg, '/')
180+
if p < 0 {
181+
return name, ImportSpec{Path: pkg, Name: appendpb(pkg)}, nil
182+
}
183+
elem := pkg[p+1:]
184+
if len(elem) >= 2 && elem[0] == 'v' && elem[1] >= '0' && elem[1] <= '9' {
185+
// It's a version number; skip so we get a more meaningful name
186+
pkg = pkg[:p]
187+
continue
188+
}
189+
return name, ImportSpec{Path: pkg, Name: appendpb(elem)}, nil
190+
}
191+
}
192+
193+
// ImportSpec reports the ImportSpec for package containing protobuf element e.
194+
// Deprecated: Use NameSpec instead.
195+
func (in *Info) ImportSpec(e ProtoType) (ImportSpec, error) {
196+
_, imp, err := in.NameSpec(e)
197+
return imp, err
198+
}
199+
200+
// ReduceServName removes redundant components from the service name.
201+
// For example, FooServiceV2 -> Foo.
202+
// The returned name is used as part of longer names, like FooClient.
203+
// If the package name and the service name is the same,
204+
// ReduceServName returns empty string, so we get foo.Client instead of foo.FooClient.
205+
func ReduceServName(svc, pkg string) string {
206+
// remove trailing version
207+
if p := strings.LastIndexByte(svc, 'V'); p >= 0 {
208+
isVer := true
209+
for _, r := range svc[p+1:] {
210+
if !unicode.IsDigit(r) {
211+
isVer = false
212+
break
213+
}
214+
}
215+
if isVer {
216+
svc = svc[:p]
217+
}
218+
}
219+
220+
svc = strings.TrimSuffix(svc, "Service")
221+
if strings.EqualFold(svc, pkg) {
222+
svc = ""
223+
}
224+
225+
// This is a special case for IAM and should not be
226+
// extended to support any new API name containing
227+
// an acronym.
228+
//
229+
// In order to avoid a breaking change for IAM
230+
// clients, we must keep consistent identifier casing.
231+
if strings.Contains(svc, "IAM") {
232+
svc = strings.ReplaceAll(svc, "IAM", "Iam")
233+
}
234+
235+
return svc
236+
}

internal/pbinfo/prim2go.go

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
// Copyright 2018 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package pbinfo
16+
17+
import "github.com/golang/protobuf/protoc-gen-go/descriptor"
18+
19+
// GoTypeForPrim maps protobuf primitive types to Go primitive types.
20+
var GoTypeForPrim = map[descriptor.FieldDescriptorProto_Type]string{
21+
descriptor.FieldDescriptorProto_TYPE_DOUBLE: "float64",
22+
descriptor.FieldDescriptorProto_TYPE_FLOAT: "float32",
23+
descriptor.FieldDescriptorProto_TYPE_INT64: "int64",
24+
descriptor.FieldDescriptorProto_TYPE_UINT64: "uint64",
25+
descriptor.FieldDescriptorProto_TYPE_INT32: "int32",
26+
descriptor.FieldDescriptorProto_TYPE_FIXED64: "uint64",
27+
descriptor.FieldDescriptorProto_TYPE_FIXED32: "uint32",
28+
descriptor.FieldDescriptorProto_TYPE_BOOL: "bool",
29+
descriptor.FieldDescriptorProto_TYPE_STRING: "string",
30+
descriptor.FieldDescriptorProto_TYPE_BYTES: "[]byte",
31+
descriptor.FieldDescriptorProto_TYPE_UINT32: "uint32",
32+
descriptor.FieldDescriptorProto_TYPE_SFIXED32: "int32",
33+
descriptor.FieldDescriptorProto_TYPE_SFIXED64: "int64",
34+
descriptor.FieldDescriptorProto_TYPE_SINT32: "int32",
35+
descriptor.FieldDescriptorProto_TYPE_SINT64: "int64",
36+
}

0 commit comments

Comments
 (0)