protogen.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864
  1. // Copyright 2018 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. // Package protogen provides support for writing protoc plugins.
  5. //
  6. // Plugins for protoc, the Protocol Buffers Compiler, are programs which read
  7. // a CodeGeneratorRequest protocol buffer from standard input and write a
  8. // CodeGeneratorResponse protocol buffer to standard output. This package
  9. // provides support for writing plugins which generate Go code.
  10. package protogen
  11. import (
  12. "bufio"
  13. "bytes"
  14. "fmt"
  15. "go/ast"
  16. "go/parser"
  17. "go/printer"
  18. "go/token"
  19. "io/ioutil"
  20. "os"
  21. "path"
  22. "path/filepath"
  23. "sort"
  24. "strconv"
  25. "strings"
  26. "github.com/golang/protobuf/proto"
  27. descpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
  28. pluginpb "github.com/golang/protobuf/protoc-gen-go/plugin"
  29. "github.com/golang/protobuf/v2/reflect/protoreflect"
  30. "github.com/golang/protobuf/v2/reflect/protoregistry"
  31. "github.com/golang/protobuf/v2/reflect/prototype"
  32. "golang.org/x/tools/go/ast/astutil"
  33. )
  34. // Run executes a function as a protoc plugin.
  35. //
  36. // It reads a CodeGeneratorRequest message from os.Stdin, invokes the plugin
  37. // function, and writes a CodeGeneratorResponse message to os.Stdout.
  38. //
  39. // If a failure occurs while reading or writing, Run prints an error to
  40. // os.Stderr and calls os.Exit(1).
  41. //
  42. // Passing a nil options is equivalent to passing a zero-valued one.
  43. func Run(opts *Options, f func(*Plugin) error) {
  44. if err := run(opts, f); err != nil {
  45. fmt.Fprintf(os.Stderr, "%s: %v\n", filepath.Base(os.Args[0]), err)
  46. os.Exit(1)
  47. }
  48. }
  49. func run(opts *Options, f func(*Plugin) error) error {
  50. in, err := ioutil.ReadAll(os.Stdin)
  51. if err != nil {
  52. return err
  53. }
  54. req := &pluginpb.CodeGeneratorRequest{}
  55. if err := proto.Unmarshal(in, req); err != nil {
  56. return err
  57. }
  58. gen, err := New(req, opts)
  59. if err != nil {
  60. return err
  61. }
  62. if err := f(gen); err != nil {
  63. // Errors from the plugin function are reported by setting the
  64. // error field in the CodeGeneratorResponse.
  65. //
  66. // In contrast, errors that indicate a problem in protoc
  67. // itself (unparsable input, I/O errors, etc.) are reported
  68. // to stderr.
  69. gen.Error(err)
  70. }
  71. resp := gen.Response()
  72. out, err := proto.Marshal(resp)
  73. if err != nil {
  74. return err
  75. }
  76. if _, err := os.Stdout.Write(out); err != nil {
  77. return err
  78. }
  79. return nil
  80. }
  81. // A Plugin is a protoc plugin invocation.
  82. type Plugin struct {
  83. // Request is the CodeGeneratorRequest provided by protoc.
  84. Request *pluginpb.CodeGeneratorRequest
  85. // Files is the set of files to generate and everything they import.
  86. // Files appear in topological order, so each file appears before any
  87. // file that imports it.
  88. Files []*File
  89. filesByName map[string]*File
  90. fileReg *protoregistry.Files
  91. messagesByName map[protoreflect.FullName]*Message
  92. enumsByName map[protoreflect.FullName]*Enum
  93. pathType pathType
  94. genFiles []*GeneratedFile
  95. err error
  96. }
  97. // Options are optional parameters to New.
  98. type Options struct {
  99. // If ParamFunc is non-nil, it will be called with each unknown
  100. // generator parameter.
  101. //
  102. // Plugins for protoc can accept parameters from the command line,
  103. // passed in the --<lang>_out protoc, separated from the output
  104. // directory with a colon; e.g.,
  105. //
  106. // --go_out=<param1>=<value1>,<param2>=<value2>:<output_directory>
  107. //
  108. // Parameters passed in this fashion as a comma-separated list of
  109. // key=value pairs will be passed to the ParamFunc.
  110. //
  111. // The (flag.FlagSet).Set method matches this function signature,
  112. // so parameters can be converted into flags as in the following:
  113. //
  114. // var flags flag.FlagSet
  115. // value := flags.Bool("param", false, "")
  116. // opts := &protogen.Options{
  117. // ParamFunc: flags.Set,
  118. // }
  119. // protogen.Run(opts, func(p *protogen.Plugin) error {
  120. // if *value { ... }
  121. // })
  122. ParamFunc func(name, value string) error
  123. }
  124. // New returns a new Plugin.
  125. //
  126. // Passing a nil Options is equivalent to passing a zero-valued one.
  127. func New(req *pluginpb.CodeGeneratorRequest, opts *Options) (*Plugin, error) {
  128. if opts == nil {
  129. opts = &Options{}
  130. }
  131. gen := &Plugin{
  132. Request: req,
  133. filesByName: make(map[string]*File),
  134. fileReg: protoregistry.NewFiles(),
  135. messagesByName: make(map[protoreflect.FullName]*Message),
  136. enumsByName: make(map[protoreflect.FullName]*Enum),
  137. }
  138. packageNames := make(map[string]GoPackageName) // filename -> package name
  139. importPaths := make(map[string]GoImportPath) // filename -> import path
  140. var packageImportPath GoImportPath
  141. for _, param := range strings.Split(req.GetParameter(), ",") {
  142. var value string
  143. if i := strings.Index(param, "="); i >= 0 {
  144. value = param[i+1:]
  145. param = param[0:i]
  146. }
  147. switch param {
  148. case "":
  149. // Ignore.
  150. case "import_prefix":
  151. // TODO
  152. case "import_path":
  153. packageImportPath = GoImportPath(value)
  154. case "paths":
  155. switch value {
  156. case "import":
  157. gen.pathType = pathTypeImport
  158. case "source_relative":
  159. gen.pathType = pathTypeSourceRelative
  160. default:
  161. return nil, fmt.Errorf(`unknown path type %q: want "import" or "source_relative"`, value)
  162. }
  163. case "annotate_code":
  164. // TODO
  165. default:
  166. if param[0] == 'M' {
  167. importPaths[param[1:]] = GoImportPath(value)
  168. continue
  169. }
  170. if opts.ParamFunc != nil {
  171. if err := opts.ParamFunc(param, value); err != nil {
  172. return nil, err
  173. }
  174. }
  175. }
  176. }
  177. // Figure out the import path and package name for each file.
  178. //
  179. // The rules here are complicated and have grown organically over time.
  180. // Interactions between different ways of specifying package information
  181. // may be surprising.
  182. //
  183. // The recommended approach is to include a go_package option in every
  184. // .proto source file specifying the full import path of the Go package
  185. // associated with this file.
  186. //
  187. // option go_package = "github.com/golang/protobuf/ptypes/any";
  188. //
  189. // Build systems which want to exert full control over import paths may
  190. // specify M<filename>=<import_path> flags.
  191. //
  192. // Other approaches are not recommend.
  193. generatedFileNames := make(map[string]bool)
  194. for _, name := range gen.Request.FileToGenerate {
  195. generatedFileNames[name] = true
  196. }
  197. // We need to determine the import paths before the package names,
  198. // because the Go package name for a file is sometimes derived from
  199. // different file in the same package.
  200. packageNameForImportPath := make(map[GoImportPath]GoPackageName)
  201. for _, fdesc := range gen.Request.ProtoFile {
  202. filename := fdesc.GetName()
  203. packageName, importPath := goPackageOption(fdesc)
  204. switch {
  205. case importPaths[filename] != "":
  206. // Command line: M=foo.proto=quux/bar
  207. //
  208. // Explicit mapping of source file to import path.
  209. case generatedFileNames[filename] && packageImportPath != "":
  210. // Command line: import_path=quux/bar
  211. //
  212. // The import_path flag sets the import path for every file that
  213. // we generate code for.
  214. importPaths[filename] = packageImportPath
  215. case importPath != "":
  216. // Source file: option go_package = "quux/bar";
  217. //
  218. // The go_package option sets the import path. Most users should use this.
  219. importPaths[filename] = importPath
  220. default:
  221. // Source filename.
  222. //
  223. // Last resort when nothing else is available.
  224. importPaths[filename] = GoImportPath(path.Dir(filename))
  225. }
  226. if packageName != "" {
  227. packageNameForImportPath[importPaths[filename]] = packageName
  228. }
  229. }
  230. for _, fdesc := range gen.Request.ProtoFile {
  231. filename := fdesc.GetName()
  232. packageName, _ := goPackageOption(fdesc)
  233. defaultPackageName := packageNameForImportPath[importPaths[filename]]
  234. switch {
  235. case packageName != "":
  236. // Source file: option go_package = "quux/bar";
  237. packageNames[filename] = packageName
  238. case defaultPackageName != "":
  239. // A go_package option in another file in the same package.
  240. //
  241. // This is a poor choice in general, since every source file should
  242. // contain a go_package option. Supported mainly for historical
  243. // compatibility.
  244. packageNames[filename] = defaultPackageName
  245. case generatedFileNames[filename] && packageImportPath != "":
  246. // Command line: import_path=quux/bar
  247. packageNames[filename] = cleanPackageName(path.Base(string(packageImportPath)))
  248. case fdesc.GetPackage() != "":
  249. // Source file: package quux.bar;
  250. packageNames[filename] = cleanPackageName(fdesc.GetPackage())
  251. default:
  252. // Source filename.
  253. packageNames[filename] = cleanPackageName(baseName(filename))
  254. }
  255. }
  256. // Consistency check: Every file with the same Go import path should have
  257. // the same Go package name.
  258. packageFiles := make(map[GoImportPath][]string)
  259. for filename, importPath := range importPaths {
  260. packageFiles[importPath] = append(packageFiles[importPath], filename)
  261. }
  262. for importPath, filenames := range packageFiles {
  263. for i := 1; i < len(filenames); i++ {
  264. if a, b := packageNames[filenames[0]], packageNames[filenames[i]]; a != b {
  265. return nil, fmt.Errorf("Go package %v has inconsistent names %v (%v) and %v (%v)",
  266. importPath, a, filenames[0], b, filenames[i])
  267. }
  268. }
  269. }
  270. for _, fdesc := range gen.Request.ProtoFile {
  271. filename := fdesc.GetName()
  272. if gen.filesByName[filename] != nil {
  273. return nil, fmt.Errorf("duplicate file name: %q", filename)
  274. }
  275. f, err := newFile(gen, fdesc, packageNames[filename], importPaths[filename])
  276. if err != nil {
  277. return nil, err
  278. }
  279. gen.Files = append(gen.Files, f)
  280. gen.filesByName[filename] = f
  281. }
  282. for _, filename := range gen.Request.FileToGenerate {
  283. f, ok := gen.FileByName(filename)
  284. if !ok {
  285. return nil, fmt.Errorf("no descriptor for generated file: %v", filename)
  286. }
  287. f.Generate = true
  288. }
  289. return gen, nil
  290. }
  291. // Error records an error in code generation. The generator will report the
  292. // error back to protoc and will not produce output.
  293. func (gen *Plugin) Error(err error) {
  294. if gen.err == nil {
  295. gen.err = err
  296. }
  297. }
  298. // Response returns the generator output.
  299. func (gen *Plugin) Response() *pluginpb.CodeGeneratorResponse {
  300. resp := &pluginpb.CodeGeneratorResponse{}
  301. if gen.err != nil {
  302. resp.Error = proto.String(gen.err.Error())
  303. return resp
  304. }
  305. for _, gf := range gen.genFiles {
  306. content, err := gf.Content()
  307. if err != nil {
  308. return &pluginpb.CodeGeneratorResponse{
  309. Error: proto.String(err.Error()),
  310. }
  311. }
  312. resp.File = append(resp.File, &pluginpb.CodeGeneratorResponse_File{
  313. Name: proto.String(gf.filename),
  314. Content: proto.String(string(content)),
  315. })
  316. }
  317. return resp
  318. }
  319. // FileByName returns the file with the given name.
  320. func (gen *Plugin) FileByName(name string) (f *File, ok bool) {
  321. f, ok = gen.filesByName[name]
  322. return f, ok
  323. }
  324. // A File describes a .proto source file.
  325. type File struct {
  326. Desc protoreflect.FileDescriptor
  327. Proto *descpb.FileDescriptorProto
  328. GoPackageName GoPackageName // name of this file's Go package
  329. GoImportPath GoImportPath // import path of this file's Go package
  330. Messages []*Message // top-level message declarations
  331. Enums []*Enum // top-level enum declarations
  332. Extensions []*Extension // top-level extension declarations
  333. Generate bool // true if we should generate code for this file
  334. // GeneratedFilenamePrefix is used to construct filenames for generated
  335. // files associated with this source file.
  336. //
  337. // For example, the source file "dir/foo.proto" might have a filename prefix
  338. // of "dir/foo". Appending ".pb.go" produces an output file of "dir/foo.pb.go".
  339. GeneratedFilenamePrefix string
  340. }
  341. func newFile(gen *Plugin, p *descpb.FileDescriptorProto, packageName GoPackageName, importPath GoImportPath) (*File, error) {
  342. desc, err := prototype.NewFileFromDescriptorProto(p, gen.fileReg)
  343. if err != nil {
  344. return nil, fmt.Errorf("invalid FileDescriptorProto %q: %v", p.GetName(), err)
  345. }
  346. if err := gen.fileReg.Register(desc); err != nil {
  347. return nil, fmt.Errorf("cannot register descriptor %q: %v", p.GetName(), err)
  348. }
  349. f := &File{
  350. Desc: desc,
  351. Proto: p,
  352. GoPackageName: packageName,
  353. GoImportPath: importPath,
  354. }
  355. // Determine the prefix for generated Go files.
  356. prefix := p.GetName()
  357. if ext := path.Ext(prefix); ext == ".proto" || ext == ".protodevel" {
  358. prefix = prefix[:len(prefix)-len(ext)]
  359. }
  360. if gen.pathType == pathTypeImport {
  361. // If paths=import (the default) and the file contains a go_package option
  362. // with a full import path, the output filename is derived from the Go import
  363. // path.
  364. //
  365. // Pass the paths=source_relative flag to always derive the output filename
  366. // from the input filename instead.
  367. if _, importPath := goPackageOption(p); importPath != "" {
  368. prefix = path.Join(string(importPath), path.Base(prefix))
  369. }
  370. }
  371. f.GeneratedFilenamePrefix = prefix
  372. for i, mdescs := 0, desc.Messages(); i < mdescs.Len(); i++ {
  373. f.Messages = append(f.Messages, newMessage(gen, f, nil, mdescs.Get(i)))
  374. }
  375. for i, edescs := 0, desc.Enums(); i < edescs.Len(); i++ {
  376. f.Enums = append(f.Enums, newEnum(gen, f, nil, edescs.Get(i)))
  377. }
  378. for i, extdescs := 0, desc.Extensions(); i < extdescs.Len(); i++ {
  379. f.Extensions = append(f.Extensions, newField(gen, f, nil, extdescs.Get(i)))
  380. }
  381. for _, message := range f.Messages {
  382. if err := message.init(gen); err != nil {
  383. return nil, err
  384. }
  385. }
  386. for _, extension := range f.Extensions {
  387. if err := extension.init(gen); err != nil {
  388. return nil, err
  389. }
  390. }
  391. return f, nil
  392. }
  393. // goPackageOption interprets a file's go_package option.
  394. // If there is no go_package, it returns ("", "").
  395. // If there's a simple name, it returns (pkg, "").
  396. // If the option implies an import path, it returns (pkg, impPath).
  397. func goPackageOption(d *descpb.FileDescriptorProto) (pkg GoPackageName, impPath GoImportPath) {
  398. opt := d.GetOptions().GetGoPackage()
  399. if opt == "" {
  400. return "", ""
  401. }
  402. // A semicolon-delimited suffix delimits the import path and package name.
  403. if i := strings.Index(opt, ";"); i >= 0 {
  404. return cleanPackageName(opt[i+1:]), GoImportPath(opt[:i])
  405. }
  406. // The presence of a slash implies there's an import path.
  407. if i := strings.LastIndex(opt, "/"); i >= 0 {
  408. return cleanPackageName(opt[i+1:]), GoImportPath(opt)
  409. }
  410. return cleanPackageName(opt), ""
  411. }
  412. // A Message describes a message.
  413. type Message struct {
  414. Desc protoreflect.MessageDescriptor
  415. GoIdent GoIdent // name of the generated Go type
  416. Fields []*Field // message field declarations
  417. Oneofs []*Oneof // oneof declarations
  418. Messages []*Message // nested message declarations
  419. Enums []*Enum // nested enum declarations
  420. Extensions []*Extension // nested extension declarations
  421. Path []int32 // location path of this message
  422. }
  423. func newMessage(gen *Plugin, f *File, parent *Message, desc protoreflect.MessageDescriptor) *Message {
  424. var path []int32
  425. if parent != nil {
  426. path = pathAppend(parent.Path, messageMessageField, int32(desc.Index()))
  427. } else {
  428. path = []int32{fileMessageField, int32(desc.Index())}
  429. }
  430. message := &Message{
  431. Desc: desc,
  432. GoIdent: newGoIdent(f, desc),
  433. Path: path,
  434. }
  435. gen.messagesByName[desc.FullName()] = message
  436. for i, mdescs := 0, desc.Messages(); i < mdescs.Len(); i++ {
  437. message.Messages = append(message.Messages, newMessage(gen, f, message, mdescs.Get(i)))
  438. }
  439. for i, edescs := 0, desc.Enums(); i < edescs.Len(); i++ {
  440. message.Enums = append(message.Enums, newEnum(gen, f, message, edescs.Get(i)))
  441. }
  442. for i, odescs := 0, desc.Oneofs(); i < odescs.Len(); i++ {
  443. message.Oneofs = append(message.Oneofs, newOneof(gen, f, message, odescs.Get(i)))
  444. }
  445. for i, fdescs := 0, desc.Fields(); i < fdescs.Len(); i++ {
  446. message.Fields = append(message.Fields, newField(gen, f, message, fdescs.Get(i)))
  447. }
  448. for i, extdescs := 0, desc.Extensions(); i < extdescs.Len(); i++ {
  449. message.Extensions = append(message.Extensions, newField(gen, f, message, extdescs.Get(i)))
  450. }
  451. // Field name conflict resolution.
  452. //
  453. // We assume well-known method names that may be attached to a generated
  454. // message type, as well as a 'Get*' method for each field. For each
  455. // field in turn, we add _s to its name until there are no conflicts.
  456. //
  457. // Any change to the following set of method names is a potential
  458. // incompatible API change because it may change generated field names.
  459. //
  460. // TODO: If we ever support a 'go_name' option to set the Go name of a
  461. // field, we should consider dropping this entirely. The conflict
  462. // resolution algorithm is subtle and surprising (changing the order
  463. // in which fields appear in the .proto source file can change the
  464. // names of fields in generated code), and does not adapt well to
  465. // adding new per-field methods such as setters.
  466. usedNames := map[string]bool{
  467. "Reset": true,
  468. "String": true,
  469. "ProtoMessage": true,
  470. "Marshal": true,
  471. "Unmarshal": true,
  472. "ExtensionRangeArray": true,
  473. "ExtensionMap": true,
  474. "Descriptor": true,
  475. }
  476. makeNameUnique := func(name string) string {
  477. for usedNames[name] || usedNames["Get"+name] {
  478. name += "_"
  479. }
  480. usedNames[name] = true
  481. usedNames["Get"+name] = true
  482. return name
  483. }
  484. seenOneofs := make(map[int]bool)
  485. for _, field := range message.Fields {
  486. field.GoName = makeNameUnique(field.GoName)
  487. if field.OneofType != nil {
  488. if !seenOneofs[field.OneofType.Desc.Index()] {
  489. // If this is a field in a oneof that we haven't seen before,
  490. // make the name for that oneof unique as well.
  491. field.OneofType.GoName = makeNameUnique(field.OneofType.GoName)
  492. seenOneofs[field.OneofType.Desc.Index()] = true
  493. }
  494. }
  495. }
  496. return message
  497. }
  498. func (message *Message) init(gen *Plugin) error {
  499. for _, child := range message.Messages {
  500. if err := child.init(gen); err != nil {
  501. return err
  502. }
  503. }
  504. for _, field := range message.Fields {
  505. if err := field.init(gen); err != nil {
  506. return err
  507. }
  508. }
  509. for _, oneof := range message.Oneofs {
  510. oneof.init(gen, message)
  511. }
  512. for _, extension := range message.Extensions {
  513. if err := extension.init(gen); err != nil {
  514. return err
  515. }
  516. }
  517. return nil
  518. }
  519. // A Field describes a message field.
  520. type Field struct {
  521. Desc protoreflect.FieldDescriptor
  522. // GoName is the base name of this field's Go field and methods.
  523. // For code generated by protoc-gen-go, this means a field named
  524. // '{{GoName}}' and a getter method named 'Get{{GoName}}'.
  525. GoName string
  526. ParentMessage *Message // message in which this field is defined; nil if top-level extension
  527. ExtendedType *Message // extended message for extension fields; nil otherwise
  528. MessageType *Message // type for message or group fields; nil otherwise
  529. EnumType *Enum // type for enum fields; nil otherwise
  530. OneofType *Oneof // containing oneof; nil if not part of a oneof
  531. Path []int32 // location path of this field
  532. }
  533. func newField(gen *Plugin, f *File, message *Message, desc protoreflect.FieldDescriptor) *Field {
  534. var path []int32
  535. switch {
  536. case desc.ExtendedType() != nil && message == nil:
  537. path = []int32{fileExtensionField, int32(desc.Index())}
  538. case desc.ExtendedType() != nil && message != nil:
  539. path = pathAppend(message.Path, messageExtensionField, int32(desc.Index()))
  540. default:
  541. path = pathAppend(message.Path, messageFieldField, int32(desc.Index()))
  542. }
  543. field := &Field{
  544. Desc: desc,
  545. GoName: camelCase(string(desc.Name())),
  546. ParentMessage: message,
  547. Path: path,
  548. }
  549. if desc.OneofType() != nil {
  550. field.OneofType = message.Oneofs[desc.OneofType().Index()]
  551. }
  552. return field
  553. }
  554. // Extension is an alias of Field for documentation.
  555. type Extension = Field
  556. func (field *Field) init(gen *Plugin) error {
  557. desc := field.Desc
  558. switch desc.Kind() {
  559. case protoreflect.MessageKind, protoreflect.GroupKind:
  560. mname := desc.MessageType().FullName()
  561. message, ok := gen.messagesByName[mname]
  562. if !ok {
  563. return fmt.Errorf("field %v: no descriptor for type %v", desc.FullName(), mname)
  564. }
  565. field.MessageType = message
  566. case protoreflect.EnumKind:
  567. ename := field.Desc.EnumType().FullName()
  568. enum, ok := gen.enumsByName[ename]
  569. if !ok {
  570. return fmt.Errorf("field %v: no descriptor for enum %v", desc.FullName(), ename)
  571. }
  572. field.EnumType = enum
  573. }
  574. if desc.ExtendedType() != nil {
  575. mname := desc.ExtendedType().FullName()
  576. message, ok := gen.messagesByName[mname]
  577. if !ok {
  578. return fmt.Errorf("field %v: no descriptor for type %v", desc.FullName(), mname)
  579. }
  580. field.ExtendedType = message
  581. }
  582. return nil
  583. }
  584. // A Oneof describes a oneof field.
  585. type Oneof struct {
  586. Desc protoreflect.OneofDescriptor
  587. GoName string // Go field name of this oneof
  588. ParentMessage *Message // message in which this oneof occurs
  589. Fields []*Field // fields that are part of this oneof
  590. Path []int32 // location path of this oneof
  591. }
  592. func newOneof(gen *Plugin, f *File, message *Message, desc protoreflect.OneofDescriptor) *Oneof {
  593. return &Oneof{
  594. Desc: desc,
  595. ParentMessage: message,
  596. GoName: camelCase(string(desc.Name())),
  597. Path: pathAppend(message.Path, messageOneofField, int32(desc.Index())),
  598. }
  599. }
  600. func (oneof *Oneof) init(gen *Plugin, parent *Message) {
  601. for i, fdescs := 0, oneof.Desc.Fields(); i < fdescs.Len(); i++ {
  602. oneof.Fields = append(oneof.Fields, parent.Fields[fdescs.Get(i).Index()])
  603. }
  604. }
  605. // An Enum describes an enum.
  606. type Enum struct {
  607. Desc protoreflect.EnumDescriptor
  608. GoIdent GoIdent // name of the generated Go type
  609. Values []*EnumValue // enum values
  610. Path []int32 // location path of this enum
  611. }
  612. func newEnum(gen *Plugin, f *File, parent *Message, desc protoreflect.EnumDescriptor) *Enum {
  613. var path []int32
  614. if parent != nil {
  615. path = pathAppend(parent.Path, messageEnumField, int32(desc.Index()))
  616. } else {
  617. path = []int32{fileEnumField, int32(desc.Index())}
  618. }
  619. enum := &Enum{
  620. Desc: desc,
  621. GoIdent: newGoIdent(f, desc),
  622. Path: path,
  623. }
  624. gen.enumsByName[desc.FullName()] = enum
  625. for i, evdescs := 0, enum.Desc.Values(); i < evdescs.Len(); i++ {
  626. enum.Values = append(enum.Values, newEnumValue(gen, f, parent, enum, evdescs.Get(i)))
  627. }
  628. return enum
  629. }
  630. // An EnumValue describes an enum value.
  631. type EnumValue struct {
  632. Desc protoreflect.EnumValueDescriptor
  633. GoIdent GoIdent // name of the generated Go type
  634. Path []int32 // location path of this enum value
  635. }
  636. func newEnumValue(gen *Plugin, f *File, message *Message, enum *Enum, desc protoreflect.EnumValueDescriptor) *EnumValue {
  637. // A top-level enum value's name is: EnumName_ValueName
  638. // An enum value contained in a message is: MessageName_ValueName
  639. //
  640. // Enum value names are not camelcased.
  641. parentIdent := enum.GoIdent
  642. if message != nil {
  643. parentIdent = message.GoIdent
  644. }
  645. name := parentIdent.GoName + "_" + string(desc.Name())
  646. return &EnumValue{
  647. Desc: desc,
  648. GoIdent: GoIdent{
  649. GoName: name,
  650. GoImportPath: f.GoImportPath,
  651. },
  652. Path: pathAppend(enum.Path, enumValueField, int32(desc.Index())),
  653. }
  654. }
  655. // A GeneratedFile is a generated file.
  656. type GeneratedFile struct {
  657. filename string
  658. goImportPath GoImportPath
  659. buf bytes.Buffer
  660. packageNames map[GoImportPath]GoPackageName
  661. usedPackageNames map[GoPackageName]bool
  662. manualImports map[GoImportPath]bool
  663. }
  664. // NewGeneratedFile creates a new generated file with the given filename
  665. // and import path.
  666. func (gen *Plugin) NewGeneratedFile(filename string, goImportPath GoImportPath) *GeneratedFile {
  667. g := &GeneratedFile{
  668. filename: filename,
  669. goImportPath: goImportPath,
  670. packageNames: make(map[GoImportPath]GoPackageName),
  671. usedPackageNames: make(map[GoPackageName]bool),
  672. manualImports: make(map[GoImportPath]bool),
  673. }
  674. gen.genFiles = append(gen.genFiles, g)
  675. return g
  676. }
  677. // P prints a line to the generated output. It converts each parameter to a
  678. // string following the same rules as fmt.Print. It never inserts spaces
  679. // between parameters.
  680. //
  681. // TODO: .meta file annotations.
  682. func (g *GeneratedFile) P(v ...interface{}) {
  683. for _, x := range v {
  684. switch x := x.(type) {
  685. case GoIdent:
  686. fmt.Fprint(&g.buf, g.QualifiedGoIdent(x))
  687. default:
  688. fmt.Fprint(&g.buf, x)
  689. }
  690. }
  691. fmt.Fprintln(&g.buf)
  692. }
  693. // QualifiedGoIdent returns the string to use for a Go identifier.
  694. //
  695. // If the identifier is from a different Go package than the generated file,
  696. // the returned name will be qualified (package.name) and an import statement
  697. // for the identifier's package will be included in the file.
  698. func (g *GeneratedFile) QualifiedGoIdent(ident GoIdent) string {
  699. if ident.GoImportPath == g.goImportPath {
  700. return ident.GoName
  701. }
  702. if packageName, ok := g.packageNames[ident.GoImportPath]; ok {
  703. return string(packageName) + "." + ident.GoName
  704. }
  705. packageName := cleanPackageName(baseName(string(ident.GoImportPath)))
  706. for i, orig := 1, packageName; g.usedPackageNames[packageName]; i++ {
  707. packageName = orig + GoPackageName(strconv.Itoa(i))
  708. }
  709. g.packageNames[ident.GoImportPath] = packageName
  710. g.usedPackageNames[packageName] = true
  711. return string(packageName) + "." + ident.GoName
  712. }
  713. // Import ensures a package is imported by the generated file.
  714. //
  715. // Packages referenced by QualifiedGoIdent are automatically imported.
  716. // Explicitly importing a package with Import is generally only necessary
  717. // when the import will be blank (import _ "package").
  718. func (g *GeneratedFile) Import(importPath GoImportPath) {
  719. g.manualImports[importPath] = true
  720. }
  721. // Write implements io.Writer.
  722. func (g *GeneratedFile) Write(p []byte) (n int, err error) {
  723. return g.buf.Write(p)
  724. }
  725. // Content returns the contents of the generated file.
  726. func (g *GeneratedFile) Content() ([]byte, error) {
  727. if !strings.HasSuffix(g.filename, ".go") {
  728. return g.buf.Bytes(), nil
  729. }
  730. // Reformat generated code.
  731. original := g.buf.Bytes()
  732. fset := token.NewFileSet()
  733. file, err := parser.ParseFile(fset, "", original, parser.ParseComments)
  734. if err != nil {
  735. // Print out the bad code with line numbers.
  736. // This should never happen in practice, but it can while changing generated code
  737. // so consider this a debugging aid.
  738. var src bytes.Buffer
  739. s := bufio.NewScanner(bytes.NewReader(original))
  740. for line := 1; s.Scan(); line++ {
  741. fmt.Fprintf(&src, "%5d\t%s\n", line, s.Bytes())
  742. }
  743. return nil, fmt.Errorf("%v: unparsable Go source: %v\n%v", g.filename, err, src.String())
  744. }
  745. // Add imports.
  746. var importPaths []string
  747. for importPath := range g.packageNames {
  748. importPaths = append(importPaths, string(importPath))
  749. }
  750. sort.Strings(importPaths)
  751. for _, importPath := range importPaths {
  752. astutil.AddNamedImport(fset, file, string(g.packageNames[GoImportPath(importPath)]), importPath)
  753. }
  754. for importPath := range g.manualImports {
  755. if _, ok := g.packageNames[importPath]; ok {
  756. continue
  757. }
  758. astutil.AddNamedImport(fset, file, "_", string(importPath))
  759. }
  760. ast.SortImports(fset, file)
  761. var out bytes.Buffer
  762. if err = (&printer.Config{Mode: printer.TabIndent | printer.UseSpaces, Tabwidth: 8}).Fprint(&out, fset, file); err != nil {
  763. return nil, fmt.Errorf("%v: can not reformat Go source: %v", g.filename, err)
  764. }
  765. // TODO: Annotations.
  766. return out.Bytes(), nil
  767. }
  768. type pathType int
  769. const (
  770. pathTypeImport pathType = iota
  771. pathTypeSourceRelative
  772. )
  773. // The SourceCodeInfo message describes the location of elements of a parsed
  774. // .proto file by way of a "path", which is a sequence of integers that
  775. // describe the route from a FileDescriptorProto to the relevant submessage.
  776. // The path alternates between a field number of a repeated field, and an index
  777. // into that repeated field. The constants below define the field numbers that
  778. // are used.
  779. //
  780. // See descriptor.proto for more information about this.
  781. const (
  782. // field numbers in FileDescriptorProto
  783. filePackageField = 2 // package
  784. fileMessageField = 4 // message_type
  785. fileEnumField = 5 // enum_type
  786. fileExtensionField = 7 // extension
  787. // field numbers in DescriptorProto
  788. messageFieldField = 2 // field
  789. messageMessageField = 3 // nested_type
  790. messageEnumField = 4 // enum_type
  791. messageExtensionField = 6 // extension
  792. messageOneofField = 8 // oneof_decl
  793. // field numbers in EnumDescriptorProto
  794. enumValueField = 2 // value
  795. )
  796. // pathAppend appends elements to a location path.
  797. // It does not alias the original path.
  798. func pathAppend(path []int32, a ...int32) []int32 {
  799. var n []int32
  800. n = append(n, path...)
  801. n = append(n, a...)
  802. return n
  803. }