You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

119 lines
3.4 KiB

  1. // msgp is a code generation tool for
  2. // creating methods to serialize and de-serialize
  3. // Go data structures to and from MessagePack.
  4. //
  5. // This package is targeted at the `go generate` tool.
  6. // To use it, include the following directive in a
  7. // go source file with types requiring source generation:
  8. //
  9. // //go:generate msgp
  10. //
  11. // The go generate tool should set the proper environment variables for
  12. // the generator to execute without any command-line flags. However, the
  13. // following options are supported, if you need them:
  14. //
  15. // -o = output file name (default is {input}_gen.go)
  16. // -file = input file name (or directory; default is $GOFILE, which is set by the `go generate` command)
  17. // -io = satisfy the `msgp.Decodable` and `msgp.Encodable` interfaces (default is true)
  18. // -marshal = satisfy the `msgp.Marshaler` and `msgp.Unmarshaler` interfaces (default is true)
  19. // -tests = generate tests and benchmarks (default is true)
  20. //
  21. // For more information, please read README.md, and the wiki at github.com/tinylib/msgp
  22. //
  23. package main
  24. import (
  25. "flag"
  26. "fmt"
  27. "os"
  28. "path/filepath"
  29. "strings"
  30. "github.com/tinylib/msgp/gen"
  31. "github.com/tinylib/msgp/parse"
  32. "github.com/tinylib/msgp/printer"
  33. "github.com/ttacon/chalk"
  34. )
  35. var (
  36. out = flag.String("o", "", "output file")
  37. file = flag.String("file", "", "input file")
  38. encode = flag.Bool("io", true, "create Encode and Decode methods")
  39. marshal = flag.Bool("marshal", true, "create Marshal and Unmarshal methods")
  40. tests = flag.Bool("tests", true, "create tests and benchmarks")
  41. unexported = flag.Bool("unexported", false, "also process unexported types")
  42. )
  43. func main() {
  44. flag.Parse()
  45. // GOFILE is set by go generate
  46. if *file == "" {
  47. *file = os.Getenv("GOFILE")
  48. if *file == "" {
  49. fmt.Println(chalk.Red.Color("No file to parse."))
  50. os.Exit(1)
  51. }
  52. }
  53. var mode gen.Method
  54. if *encode {
  55. mode |= (gen.Encode | gen.Decode | gen.Size)
  56. }
  57. if *marshal {
  58. mode |= (gen.Marshal | gen.Unmarshal | gen.Size)
  59. }
  60. if *tests {
  61. mode |= gen.Test
  62. }
  63. if mode&^gen.Test == 0 {
  64. fmt.Println(chalk.Red.Color("No methods to generate; -io=false && -marshal=false"))
  65. os.Exit(1)
  66. }
  67. if err := Run(*file, mode, *unexported); err != nil {
  68. fmt.Println(chalk.Red.Color(err.Error()))
  69. os.Exit(1)
  70. }
  71. }
  72. // Run writes all methods using the associated file or path, e.g.
  73. //
  74. // err := msgp.Run("path/to/myfile.go", gen.Size|gen.Marshal|gen.Unmarshal|gen.Test, false)
  75. //
  76. func Run(gofile string, mode gen.Method, unexported bool) error {
  77. if mode&^gen.Test == 0 {
  78. return nil
  79. }
  80. fmt.Println(chalk.Magenta.Color("======== MessagePack Code Generator ======="))
  81. fmt.Printf(chalk.Magenta.Color(">>> Input: \"%s\"\n"), gofile)
  82. fs, err := parse.File(gofile, unexported)
  83. if err != nil {
  84. return err
  85. }
  86. if len(fs.Identities) == 0 {
  87. fmt.Println(chalk.Magenta.Color("No types requiring code generation were found!"))
  88. return nil
  89. }
  90. return printer.PrintFile(newFilename(gofile, fs.Package), fs, mode)
  91. }
  92. // picks a new file name based on input flags and input filename(s).
  93. func newFilename(old string, pkg string) string {
  94. if *out != "" {
  95. if pre := strings.TrimPrefix(*out, old); len(pre) > 0 &&
  96. !strings.HasSuffix(*out, ".go") {
  97. return filepath.Join(old, *out)
  98. }
  99. return *out
  100. }
  101. if fi, err := os.Stat(old); err == nil && fi.IsDir() {
  102. old = filepath.Join(old, pkg)
  103. }
  104. // new file name is old file name + _gen.go
  105. return strings.TrimSuffix(old, ".go") + "_gen.go"
  106. }