/
generator.go
59 lines (47 loc) · 1.18 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
package internal
import (
"go/format"
"google.golang.org/protobuf/compiler/protogen"
)
type generator struct {
}
func New() *generator {
return &generator{}
}
// generateFile generates a _nats.pb.go file containing nats service definitions.
func (g *generator) GenerateFile(gen *protogen.Plugin, file *protogen.File) (*protogen.GeneratedFile, error) {
filename := file.GeneratedFilenamePrefix + "_nats.pb.go"
gf := gen.NewGeneratedFile(filename, file.GoImportPath)
gf.P("// Code generated by protoc-gen-go-nats. DO NOT EDIT.")
gf.P()
gf.P("package ", file.GoPackageName)
gf.P()
gf.P(`
import (
"context"
"time"
nats "github.com/nats-io/nats.go"
natspb "github.com/jonashiltl/nats.pb"
micro "github.com/nats-io/nats.go/micro"
proto "google.golang.org/protobuf/proto"
protodesc "google.golang.org/protobuf/reflect/protodesc"
)
`)
for _, srv := range file.Services {
code, err := applyTemplate(applyTemplateParams{
srv: srv,
})
if err != nil {
return nil, err
}
formatted, err := format.Source([]byte(code))
if err != nil {
return nil, err
}
_, err = gf.Write(formatted)
if err != nil {
return nil, err
}
}
return gf, nil
}