forked from moul/protoc-gen-gotemplate
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
148 lines (129 loc) · 3.51 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package main
import (
"io/ioutil"
"log"
"os"
"strings"
"github.com/golang/protobuf/proto"
"github.com/golang/protobuf/protoc-gen-go/generator"
"github.com/golang/protobuf/protoc-gen-go/plugin"
ggdescriptor "github.com/grpc-ecosystem/grpc-gateway/protoc-gen-grpc-gateway/descriptor"
pgghelpers "github.com/moul/protoc-gen-gotemplate/helpers"
)
var (
registry *ggdescriptor.Registry // some helpers need access to registry
)
const (
boolTrue = "true"
boolFalse = "false"
)
func main() {
g := generator.New()
data, err := ioutil.ReadAll(os.Stdin)
if err != nil {
g.Error(err, "reading input")
}
if err = proto.Unmarshal(data, g.Request); err != nil {
g.Error(err, "parsing input proto")
}
if len(g.Request.FileToGenerate) == 0 {
g.Fail("no files to generate")
}
g.CommandLineParameters(g.Request.GetParameter())
// Parse parameters
var (
templateDir = "./templates"
destinationDir = "."
debug = false
all = false
singlePackageMode = false
)
if parameter := g.Request.GetParameter(); parameter != "" {
for _, param := range strings.Split(parameter, ",") {
parts := strings.Split(param, "=")
if len(parts) != 2 {
log.Printf("Err: invalid parameter: %q", param)
continue
}
switch parts[0] {
case "template_dir":
templateDir = parts[1]
case "destination_dir":
destinationDir = parts[1]
case "single-package-mode":
switch strings.ToLower(parts[1]) {
case boolTrue, "t":
singlePackageMode = true
case boolFalse, "f":
default:
log.Printf("Err: invalid value for single-package-mode: %q", parts[1])
}
case "debug":
switch strings.ToLower(parts[1]) {
case boolTrue, "t":
debug = true
case boolFalse, "f":
default:
log.Printf("Err: invalid value for debug: %q", parts[1])
}
case "all":
switch strings.ToLower(parts[1]) {
case boolTrue, "t":
all = true
case boolFalse, "f":
default:
log.Printf("Err: invalid value for debug: %q", parts[1])
}
default:
log.Printf("Err: unknown parameter: %q", param)
}
}
}
tmplMap := make(map[string]*plugin_go.CodeGeneratorResponse_File)
concatOrAppend := func(file *plugin_go.CodeGeneratorResponse_File) {
if val, ok := tmplMap[file.GetName()]; ok {
*val.Content += file.GetContent()
} else {
tmplMap[file.GetName()] = file
g.Response.File = append(g.Response.File, file)
}
}
if singlePackageMode {
registry = ggdescriptor.NewRegistry()
pgghelpers.SetRegistry(registry)
if err = registry.Load(g.Request); err != nil {
g.Error(err, "registry: failed to load the request")
}
}
// Generate the encoders
for _, file := range g.Request.GetProtoFile() {
if all {
if singlePackageMode {
if _, err = registry.LookupFile(file.GetName()); err != nil {
g.Error(err, "registry: failed to lookup file %q", file.GetName())
}
}
encoder := NewGenericTemplateBasedEncoder(templateDir, file, debug, destinationDir)
for _, tmpl := range encoder.Files() {
concatOrAppend(tmpl)
}
continue
}
for _, service := range file.GetService() {
encoder := NewGenericServiceTemplateBasedEncoder(templateDir, service, file, debug, destinationDir)
for _, tmpl := range encoder.Files() {
concatOrAppend(tmpl)
}
}
}
// Generate the protobufs
g.GenerateAllFiles()
data, err = proto.Marshal(g.Response)
if err != nil {
g.Error(err, "failed to marshal output proto")
}
_, err = os.Stdout.Write(data)
if err != nil {
g.Error(err, "failed to write output proto")
}
}