forked from lyft/protoc-gen-star
/
generator.go
134 lines (110 loc) · 4.33 KB
/
generator.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
package pgs
import (
"io"
"log"
"os"
"github.com/golang/protobuf/protoc-gen-go/generator"
)
// Generator replaces the standard protoc-gen-go generator.Generator. It
// permits the registration of both standard protoc-gen-go plugins (eg, grpc)
// that are in-band with the officially generated code as well as Modules which
// enable creating out-of-band code using a computed protobuf AST.
type Generator struct {
Debugger
pgg ProtocGenGo // protoc-gen-go generator
gatherer *gatherer // gatherer pgg plugin
persister persister // handles writing artifacts to their output
workflow workflow // handles the actual code generation execution
plugins []Plugin // registered pgg plugins
mods []Module // registered pg* modules
in io.Reader // protoc input reader
out io.Writer // protoc output writer
debug bool // whether or not to print debug messages
includeGo bool // whether or not to gen official go code
params Parameters // CLI parameters passed in from protoc
paramMutators []ParamMutator // registered param mutators
}
// Init configures a new Generator. InitOptions may be provided as well to
// modify the behavior of the generator.
func Init(opts ...InitOption) *Generator {
g := &Generator{
pgg: Wrap(generator.New()),
gatherer: newGatherer(),
in: os.Stdin,
out: os.Stdout,
persister: newPersister(),
workflow: new(standardWorkflow),
}
g.persister.SetPGG(g.pgg)
for _, opt := range opts {
opt(g)
}
g.Debugger = initDebugger(g, log.New(os.Stderr, "", 0))
g.persister.SetDebugger(g.Debugger)
if !g.includeGo {
g.workflow = &excludeGoWorkflow{workflow: g.workflow}
}
g.workflow = &onceWorkflow{workflow: g.workflow}
return g
}
// RegisterPlugin attaches protoc-gen-go plugins to the Generator. If p
// implements the protoc-gen-star Plugin interface, a Debugger will be passed
// in. This method is solely a wrapper around generator.RegisterPlugin. When
// designing these, all context should be cleared when Init is called. Note
// that these are currently global in scope and not specific to this generator
// instance.
func (g *Generator) RegisterPlugin(p ...generator.Plugin) *Generator {
for _, pl := range p {
g.Assert(pl != nil, "nil plugin provided")
g.Debug("registering plugin:", pl.Name())
if ppl, ok := pl.(Plugin); ok {
g.plugins = append(g.plugins, ppl)
}
generator.RegisterPlugin(pl)
}
return g
}
// RegisterModule should be called after Init but before Render to attach a
// custom Module to the Generator. This method can be called multiple times.
func (g *Generator) RegisterModule(m ...Module) *Generator {
for _, mod := range m {
g.Assert(mod != nil, "nil module provided")
g.Debug("registering module: ", mod.Name())
}
g.mods = append(g.mods, m...)
return g
}
// RegisterPostProcessor should be called after Init but before Render to
// attach PostProcessors to the Generator. This method can be called multiple
// times. PostProcessors are executed against their matches in the order in
// which they are registered. Only Artifacts generated by Modules are processed.
func (g *Generator) RegisterPostProcessor(p ...PostProcessor) *Generator {
for _, pp := range p {
g.Assert(pp != nil, "nil post-processor provided")
}
g.persister.AddPostProcessor(p...)
return g
}
// AST returns the target Packages as well as all loaded Packages from the
// gatherer. Calling this method will trigger running the underlying
// protoc-gen-go workflow, including any registered plugins. Render can be
// safely called before or after this method without impacting module execution
// or writing the output to the output io.Writer. This method is particularly
// useful for integration-type tests.
func (g *Generator) AST() (targets map[string]Package, pkgs map[string]Package) {
g.workflow.Init(g)
g.workflow.Go()
return g.gatherer.targets, g.gatherer.pkgs
}
// Render emits all generated files from the plugins and Modules to the output
// io.Writer. If out is nil, os.Stdout is used. Render can only be called once
// and should be preceded by Init and any number of RegisterPlugin and/or
// RegisterModule calls.
func (g *Generator) Render() {
g.workflow.Init(g)
g.workflow.Go()
g.workflow.Star()
g.workflow.Persist()
}
func (g *Generator) push(prefix string) { g.Debugger = g.Push(prefix) }
func (g *Generator) pop() { g.Debugger = g.Pop() }