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
|
package main
import (
"fmt"
"io"
"strings"
"bytes"
pgs "github.com/lyft/protoc-gen-star/v2"
)
type PrinterModule struct {
*pgs.ModuleBase
}
func ASTPrinter() *PrinterModule { return &PrinterModule{ModuleBase: &pgs.ModuleBase{}} }
func (p *PrinterModule) Name() string { return "printer" }
func (p *PrinterModule) Execute(targets map[string]pgs.File, packages map[string]pgs.Package) []pgs.Artifact {
buf := &bytes.Buffer{}
for _, f := range targets {
p.printFile(f, buf)
}
return p.Artifacts()
}
func (p *PrinterModule) printFile(f pgs.File, buf *bytes.Buffer) {
p.Push(f.Name().String())
defer p.Pop()
buf.Reset()
v := initPrintVisitor(buf, "")
p.CheckErr(pgs.Walk(v, f), "unable to print AST tree")
out := buf.String()
if ok, _ := p.Parameters().Bool("log_tree"); ok {
p.Logf("Proto Tree:\n%s", out)
}
p.AddGeneratorFile(
f.InputPath().SetExt(".tree.txt").String(),
out,
)
}
const (
startNodePrefix = "┳ "
subNodePrefix = "┃"
leafNodePrefix = "┣"
leafNodeSpacer = "━ "
)
type PrinterVisitor struct {
pgs.Visitor
prefix string
w io.Writer
}
func initPrintVisitor(w io.Writer, prefix string) pgs.Visitor {
v := PrinterVisitor{
prefix: prefix,
w: w,
}
v.Visitor = pgs.PassThroughVisitor(&v)
return v
}
func (v PrinterVisitor) leafPrefix() string {
if strings.HasSuffix(v.prefix, subNodePrefix) {
return strings.TrimSuffix(v.prefix, subNodePrefix) + leafNodePrefix
}
return v.prefix
}
func (v PrinterVisitor) writeSubNode(str string) pgs.Visitor {
fmt.Fprintf(v.w, "%s%s%s\n", v.leafPrefix(), startNodePrefix, str)
return initPrintVisitor(v.w, fmt.Sprintf("%s%v", v.prefix, subNodePrefix))
}
func (v PrinterVisitor) writeLeaf(str string) {
fmt.Fprintf(v.w, "%s%s%s\n", v.leafPrefix(), leafNodeSpacer, str)
}
func (v PrinterVisitor) VisitFile(f pgs.File) (pgs.Visitor, error) {
return v.writeSubNode("File: " + f.Name().String()), nil
}
func (v PrinterVisitor) VisitMessage(m pgs.Message) (pgs.Visitor, error) {
return v.writeSubNode("Message: " + m.Name().String()), nil
}
func (v PrinterVisitor) VisitEnum(e pgs.Enum) (pgs.Visitor, error) {
return v.writeSubNode("Enum: " + e.Name().String()), nil
}
func (v PrinterVisitor) VisitService(s pgs.Service) (pgs.Visitor, error) {
return v.writeSubNode("Service: " + s.Name().String()), nil
}
func (v PrinterVisitor) VisitEnumValue(ev pgs.EnumValue) (pgs.Visitor, error) {
v.writeLeaf(ev.Name().String())
return nil, nil
}
func (v PrinterVisitor) VisitField(f pgs.Field) (pgs.Visitor, error) {
v.writeLeaf(f.Name().String())
return nil, nil
}
func (v PrinterVisitor) VisitMethod(m pgs.Method) (pgs.Visitor, error) {
v.writeLeaf(m.Name().String())
return nil, nil
}
|
