@@ -17,13 +17,14 @@ limitations under the License.
17
17
package genall
18
18
19
19
import (
20
+ "encoding/json"
20
21
"fmt"
21
22
"io"
22
23
"io/ioutil"
23
24
"os"
24
25
25
26
"golang.org/x/tools/go/packages"
26
- "sigs.k8s.io /yaml"
27
+ rawyaml "gopkg.in /yaml.v2 "
27
28
28
29
"sigs.k8s.io/controller-tools/pkg/loader"
29
30
"sigs.k8s.io/controller-tools/pkg/markers"
@@ -120,18 +121,30 @@ type GenerationContext struct {
120
121
InputRule
121
122
}
122
123
124
+ // WriteYAMLOptions implements the Options Pattern for WriteYAML.
125
+ type WriteYAMLOptions struct {
126
+ transform func (obj map [string ]interface {}) error
127
+ }
128
+
129
+ // WithTransform applies a transformation to objects just before writing them.
130
+ func WithTransform (transform func (obj map [string ]interface {}) error ) * WriteYAMLOptions {
131
+ return & WriteYAMLOptions {
132
+ transform : transform ,
133
+ }
134
+ }
135
+
123
136
// WriteYAML writes the given objects out, serialized as YAML, using the
124
137
// context's OutputRule. Objects are written as separate documents, separated
125
138
// from each other by `---` (as per the YAML spec).
126
- func (g GenerationContext ) WriteYAML (itemPath string , objs ... interface {}) error {
139
+ func (g GenerationContext ) WriteYAML (itemPath string , objs [] interface {}, options ... * WriteYAMLOptions ) error {
127
140
out , err := g .Open (nil , itemPath )
128
141
if err != nil {
129
142
return err
130
143
}
131
144
defer out .Close ()
132
145
133
146
for _ , obj := range objs {
134
- yamlContent , err := yaml . Marshal (obj )
147
+ yamlContent , err := yamlMarshal (obj , options ... )
135
148
if err != nil {
136
149
return err
137
150
}
@@ -147,6 +160,41 @@ func (g GenerationContext) WriteYAML(itemPath string, objs ...interface{}) error
147
160
return nil
148
161
}
149
162
163
+ // yamlMarshal is based on sigs.k8s.io/yaml.Marshal, but allows for transforming the final data before writing.
164
+ func yamlMarshal (o interface {}, options ... * WriteYAMLOptions ) ([]byte , error ) {
165
+ j , err := json .Marshal (o )
166
+ if err != nil {
167
+ return nil , fmt .Errorf ("error marshaling into JSON: %v" , err )
168
+ }
169
+
170
+ return yamlJSONToYAMLWithFilter (j , options ... )
171
+ }
172
+
173
+ // yamlJSONToYAMLWithFilter is based on sigs.k8s.io/yaml.JSONToYAML, but allows for transforming the final data before writing.
174
+ func yamlJSONToYAMLWithFilter (j []byte , options ... * WriteYAMLOptions ) ([]byte , error ) {
175
+ // Convert the JSON to an object.
176
+ var jsonObj map [string ]interface {}
177
+ // We are using yaml.Unmarshal here (instead of json.Unmarshal) because the
178
+ // Go JSON library doesn't try to pick the right number type (int, float,
179
+ // etc.) when unmarshalling to interface{}, it just picks float64
180
+ // universally. go-yaml does go through the effort of picking the right
181
+ // number type, so we can preserve number type throughout this process.
182
+ if err := rawyaml .Unmarshal (j , & jsonObj ); err != nil {
183
+ return nil , err
184
+ }
185
+
186
+ for _ , option := range options {
187
+ if option .transform != nil {
188
+ if err := option .transform (jsonObj ); err != nil {
189
+ return nil , err
190
+ }
191
+ }
192
+ }
193
+
194
+ // Marshal this object into YAML.
195
+ return rawyaml .Marshal (jsonObj )
196
+ }
197
+
150
198
// ReadFile reads the given boilerplate artifact using the context's InputRule.
151
199
func (g GenerationContext ) ReadFile (path string ) ([]byte , error ) {
152
200
file , err := g .OpenForRead (path )
0 commit comments