protogen.go 30 KB

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