Adding ability to yedit json files.

1 files changed, 19 insertions, 7 deletions

diff --git a/roles/lib_utils/library/yedit.py b/roles/lib_utils/library/yedit.py
index cf5c2e423..4bd5171a7 100644
--- a/roles/lib_utils/library/yedit.py
+++ b/roles/lib_utils/library/yedit.py
@@ -410,10 +410,16 @@ class Yedit(object):
             pass
 
         # Try to use RoundTripDumper if supported.
-        try:
-            Yedit._write(self.filename, yaml.dump(self.yaml_dict, Dumper=yaml.RoundTripDumper))
-        except AttributeError:
-            Yedit._write(self.filename, yaml.safe_dump(self.yaml_dict, default_flow_style=False))
+        if self.content_type == 'yaml':
+            try:
+                Yedit._write(self.filename, yaml.dump(self.yaml_dict, Dumper=yaml.RoundTripDumper))
+            except AttributeError:
+                Yedit._write(self.filename, yaml.safe_dump(self.yaml_dict, default_flow_style=False))
+        elif self.content_type == 'json':
+            Yedit._write(self.filename, json.dumps(self.yaml_dict, indent=4, sort_keys=True))
+        else:
+            raise YeditException('Unsupported content_type: {}.'.format(self.content_type) +
+                                 'Please specify a content_type of yaml or json.')
 
         return (True, self.yaml_dict)
 
@@ -461,7 +467,7 @@ class Yedit(object):
 
                 # Try to use RoundTripLoader if supported.
                 try:
-                    self.yaml_dict = yaml.safe_load(contents, yaml.RoundTripLoader)
+                    self.yaml_dict = yaml.load(contents, yaml.RoundTripLoader)
                 except AttributeError:
                     self.yaml_dict = yaml.safe_load(contents)
 
@@ -706,7 +712,12 @@ class Yedit(object):
 
         curr_value = invalue
         if val_type == 'yaml':
-            curr_value = yaml.load(invalue)
+            try:
+                # AUDIT:maybe-no-member makes sense due to different yaml libraries
+                # pylint: disable=maybe-no-member
+                curr_value = yaml.safe_load(invalue, Loader=yaml.RoundTripLoader)
+            except AttributeError:
+                curr_value = yaml.safe_load(invalue)
         elif val_type == 'json':
             curr_value = json.loads(invalue)
 
@@ -775,6 +786,7 @@ class Yedit(object):
         '''perform the idempotent crud operations'''
         yamlfile = Yedit(filename=params['src'],
                          backup=params['backup'],
+                         content_type=params['content_type'],
                          separator=params['separator'])
 
         state = params['state']
@@ -885,7 +897,7 @@ def main():
             debug=dict(default=False, type='bool'),
             src=dict(default=None, type='str'),
             content=dict(default=None),
-            content_type=dict(default='dict', choices=['dict']),
+            content_type=dict(default='yaml', choices=['yaml', 'json']),
             key=dict(default='', type='str'),
             value=dict(),
             value_type=dict(default='', type='str'),