/
ast.go
146 lines (128 loc) · 3.19 KB
/
ast.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
package codegen
import (
"fmt"
"go/ast"
"sort"
"strconv"
"github.com/fatih/structtag"
)
var (
_ ast.Visitor = &structFieldVisitor{}
)
// structFieldVisitor is an ast.Visitor that calls given callback function
// when visiting field in a struct.
type structFieldVisitor struct {
callback func(*ast.TypeSpec, *ast.Field)
}
func (v *structFieldVisitor) Visit(n ast.Node) ast.Visitor {
switch node := n.(type) {
case *ast.GenDecl:
for _, spec := range node.Specs {
ts, ok := spec.(*ast.TypeSpec)
if !ok {
continue
}
st, ok := ts.Type.(*ast.StructType)
if !ok {
continue
}
for _, f := range st.Fields.List {
v.callback(ts, f)
}
}
}
return v
}
// TagMap is a map of struct name -> field json name -> field tags
type TagMap map[string]map[string]*structtag.Tags
// ExtractStructFieldTagByJSONName returns an ast.Visitor that
// extracts struct field tag value by field's JSON name to tagMap.
// Kubebuilder struct field may have no name because it can be
// embedded in the struct (e.g. ObjectMeta). To cover this case,
// it depends on the field's JSON name, not the field name (which
// may be empty).
func ExtractStructFieldTagByJSONName(tagMap TagMap) ast.Visitor {
return &structFieldVisitor{
callback: func(spec *ast.TypeSpec, field *ast.Field) {
structName := spec.Name.Name
if _, ok := tagMap[structName]; !ok {
tagMap[structName] = make(map[string]*structtag.Tags)
}
tags, err := parseFieldTags(field)
if err != nil {
panic(err)
}
key, ok := getJSONName(tags)
if !ok {
return
}
tagMap[structName][key] = tags
},
}
}
// InjectStructFieldTagByJSONName returns an ast.Visitor that injects
// struct field tags by field's JSON name if the given struct field's
// tag exists in tagMap.
func InjectStructFieldTagByJSONName(tagMap TagMap) ast.Visitor {
return &structFieldVisitor{
callback: func(spec *ast.TypeSpec, field *ast.Field) {
if field.Tag == nil {
return
}
structName := spec.Name.Name
structTags, err := parseFieldTags(field)
if err != nil {
panic(err)
}
key, ok := getJSONName(structTags)
if !ok {
return
}
input, ok := tagMap[structName][key]
if !ok {
return
}
for idx := range input.Tags() {
tag := input.Tags()[idx]
// Do not override json tag
if tag.Key == "json" {
continue
}
if err := structTags.Set(tag); err != nil {
panic(err)
}
}
// Sort tags to ensure consistent output
sort.Sort(structTags)
field.Tag.Value = fmt.Sprintf("`%s`", structTags.String())
},
}
}
func parseFieldTags(field *ast.Field) (*structtag.Tags, error) {
if field.Tag == nil {
return nil, nil
}
// Remove backquotes from field tag value
rawTag, err := strconv.Unquote(field.Tag.Value)
if err != nil {
return nil, fmt.Errorf("unquote field tag value: %w", err)
}
tags, err := structtag.Parse(rawTag)
if err != nil {
return nil, fmt.Errorf("parse field tag: %w", err)
}
return tags, nil
}
func getJSONName(tags *structtag.Tags) (string, bool) {
if tags == nil {
return "", false
}
tag, err := tags.Get("json")
if err != nil {
return "", false
}
if tag.Name == "" || tag.Name == "-" {
return "", false
}
return tag.Name, true
}