forked from alvaroloes/enumer
-
Notifications
You must be signed in to change notification settings - Fork 3
/
enumer.go
197 lines (171 loc) · 5.16 KB
/
enumer.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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
package main
import "fmt"
// Arguments to format are:
// [1]: type name
const stringNameToValueMethod = `// %[1]sString retrieves an enum value from the enum constants string name.
// Throws an error if the param is not part of the enum.
func %[1]sString(s string) (%[1]s, error) {
if val, ok := _%[1]sNameToValueMap[s]; ok {
return val, nil
}
return 0, fmt.Errorf("%%s does not belong to %[1]s values", s)
}
`
// Arguments to format are:
// [1]: type name
const stringValuesMethod = `// %[1]sValues returns all values of the enum
func %[1]sValues() []%[1]s {
return _%[1]sValues
}
`
// Arguments to format are:
// [1]: type name
const stringBelongsMethodLoop = `// IsA%[1]s returns "true" if the value is listed in the enum definition. "false" otherwise
func (i %[1]s) IsA%[1]s() bool {
for _, v := range _%[1]sValues {
if i == v {
return true
}
}
return false
}
`
// Arguments to format are:
// [1]: type name
const stringBelongsMethodSet = `// IsA%[1]s returns "true" if the value is listed in the enum definition. "false" otherwise
func (i %[1]s) IsA%[1]s() bool {
_, ok := _%[1]sMap[i]
return ok
}
`
func (g *Generator) buildBasicExtras(runs [][]Value, typeName string, runsThreshold int) {
// At this moment, either "g.declareIndexAndNameVars()" or "g.declareNameVars()" has been called
// Print the slice of values
g.Printf("\nvar _%sValues = []%s{", typeName, typeName)
for _, values := range runs {
for _, value := range values {
g.Printf("\t%s, ", value.str)
}
}
g.Printf("}\n\n")
// Print the map between name and value
g.Printf("\nvar _%sNameToValueMap = map[string]%s{\n", typeName, typeName)
thereAreRuns := len(runs) > 1 && len(runs) <= runsThreshold
var n int
var runID string
for i, values := range runs {
if thereAreRuns {
runID = "_" + fmt.Sprintf("%d", i)
n = 0
} else {
runID = ""
}
for _, value := range values {
g.Printf("\t_%sName%s[%d:%d]: %s,\n", typeName, runID, n, n+len(value.name), &value)
n += len(value.name)
}
}
g.Printf("}\n\n")
// Print the basic extra methods
g.Printf(stringNameToValueMethod, typeName)
g.Printf(stringValuesMethod, typeName)
if len(runs) <= runsThreshold {
g.Printf(stringBelongsMethodLoop, typeName)
} else { // There is a map of values, the code is simpler then
g.Printf(stringBelongsMethodSet, typeName)
}
}
// Arguments to format are:
// [1]: type name
const jsonMethods = `
// MarshalJSON implements the json.Marshaler interface for %[1]s
func (i %[1]s) MarshalJSON() ([]byte, error) {
return json.Marshal(i.String())
}
// UnmarshalJSON implements the json.Unmarshaler interface for %[1]s
func (i *%[1]s) UnmarshalJSON(data []byte) error {
var s string
if err := json.Unmarshal(data, &s); err != nil {
return fmt.Errorf("%[1]s should be a string, got %%s", data)
}
var err error
*i, err = %[1]sString(s)
return err
}
`
func (g *Generator) buildJSONMethods(runs [][]Value, typeName string, runsThreshold int) {
g.Printf(jsonMethods, typeName)
}
// Arguments to format are:
// [1]: type name
const textMethods = `
// MarshalText implements the encoding.TextMarshaler interface for %[1]s
func (i %[1]s) MarshalText() ([]byte, error) {
return []byte(i.String()), nil
}
// UnmarshalText implements the encoding.TextUnmarshaler interface for %[1]s
func (i *%[1]s) UnmarshalText(text []byte) error {
var err error
*i, err = %[1]sString(string(text))
return err
}
`
func (g *Generator) buildTextMethods(runs [][]Value, typeName string, runsThreshold int) {
g.Printf(textMethods, typeName)
}
// Arguments to format are:
// [1]: type name
const yamlMethods = `
// MarshalYAML implements a YAML Marshaler for %[1]s
func (i %[1]s) MarshalYAML() (interface{}, error) {
return i.String(), nil
}
// UnmarshalYAML implements a YAML Unmarshaler for %[1]s
func (i *%[1]s) UnmarshalYAML(unmarshal func(interface{}) error) error {
var s string
if err := unmarshal(&s); err != nil {
return err
}
var err error
*i, err = %[1]sString(s)
return err
}
`
func (g *Generator) buildYAMLMethods(runs [][]Value, typeName string, runsThreshold int) {
g.Printf(yamlMethods, typeName)
}
const graphqlgoMethods = `
// ImplementsGraphQLType tells graphql-go to use this enum to resolve the %[1]s GraphQL enum type
func (%[1]s) ImplementsGraphQLType(name string) bool {
return name == "%[1]s"
}
// UnmarshalGraphQL unmarshals a GraphQL input string to the enum value.
func (i *%[1]s) UnmarshalGraphQL(input interface{}) error {
if str, ok := input.(string); ok {
if val, ok := _%[1]sNameToValueMap[str]; ok {
*i = val
return nil
}
return fmt.Errorf("%%s is not a valid %[1]s", str)
}
return fmt.Errorf("wrong type for %[1]s: %%T", input)
}
// MarshalJSON is a custom marshaler for %[1]s
//
// This function will be called whenever you
// query for fields that use the %[1]s type
func (i %[1]s) MarshalJSON() ([]byte, error) {
return json.Marshal(i.String())
}
// Maybe%[1]s returns a pointer to an enum value if the passed value is
// contained in the NameToValueMap, otherwise, it returns nil.
func Maybe%[1]s(i %[1]s) *%[1]s {
if i.IsA%[1]s() {
return &i
}
return nil
}
`
func (g *Generator) buildGraphQLGoMethods(runs [][]Value, typeName string, runsThreshold int) {
g.Printf(graphqlgoMethods, typeName)
}