1 files changed, 36 insertions, 20 deletions

diff --git a/library/modify_yaml.py b/library/modify_yaml.py
index 63b507a72..8706e80c2 100755
--- a/library/modify_yaml.py
+++ b/library/modify_yaml.py
@@ -6,6 +6,11 @@
 
 import yaml
 
+# ignore pylint errors related to the module_utils import
+# pylint: disable=redefined-builtin, unused-wildcard-import, wildcard-import
+from ansible.module_utils.basic import *  # noqa: F402,F403
+
+
 DOCUMENTATION = '''
 ---
 module: modify_yaml
@@ -21,19 +26,36 @@ EXAMPLES = '''
 '''
 
 
-# pylint: disable=missing-docstring
 def set_key(yaml_data, yaml_key, yaml_value):
+    ''' Updates a parsed yaml structure setting a key to a value.
+
+        :param yaml_data: yaml structure to modify.
+        :type yaml_data: dict
+        :param yaml_key: Key to modify.
+        :type yaml_key: mixed
+        :param yaml_value: Value use for yaml_key.
+        :type yaml_value: mixed
+        :returns: Changes to the yaml_data structure
+        :rtype: dict(tuple())
+    '''
     changes = []
     ptr = yaml_data
+    final_key = yaml_key.split('.')[-1]
     for key in yaml_key.split('.'):
-        if key not in ptr and key != yaml_key.split('.')[-1]:
+        # Key isn't present and we're not on the final key. Set to empty dictionary.
+        if key not in ptr and key != final_key:
             ptr[key] = {}
             ptr = ptr[key]
-        elif key == yaml_key.split('.')[-1]:
-            if (key in ptr and module.safe_eval(ptr[key]) != yaml_value) or (key not in ptr):
+        # Current key is the final key. Update value.
+        elif key == final_key:
+            if (key in ptr and module.safe_eval(ptr[key]) != yaml_value) or (key not in ptr):  # noqa: F405
                 ptr[key] = yaml_value
                 changes.append((yaml_key, yaml_value))
         else:
+            # Next value is None and we're not on the final key.
+            # Turn value into an empty dictionary.
+            if ptr[key] is None and key != final_key:
+                ptr[key] = {}
             ptr = ptr[key]
     return changes
 
@@ -49,7 +71,7 @@ def main():
     # redefined-outer-name
     global module
 
-    module = AnsibleModule(
+    module = AnsibleModule(  # noqa: F405
         argument_spec=dict(
             dest=dict(required=True),
             yaml_key=dict(required=True),
@@ -68,36 +90,30 @@ def main():
     # pylint: disable=missing-docstring, unused-argument
     def none_representer(dumper, data):
         return yaml.ScalarNode(tag=u'tag:yaml.org,2002:null', value=u'')
+
     yaml.add_representer(type(None), none_representer)
 
     try:
-
-        yaml_file = open(dest)
-        yaml_data = yaml.safe_load(yaml_file.read())
-        yaml_file.close()
+        with open(dest) as yaml_file:
+            yaml_data = yaml.safe_load(yaml_file.read())
 
         changes = set_key(yaml_data, yaml_key, yaml_value)
 
         if len(changes) > 0:
             if backup:
                 module.backup_local(dest)
-            yaml_file = open(dest, 'w')
-            yaml_string = yaml.dump(yaml_data, default_flow_style=False)
-            yaml_string = yaml_string.replace('\'\'', '""')
-            yaml_file.write(yaml_string)
-            yaml_file.close()
+            with open(dest, 'w') as yaml_file:
+                yaml_string = yaml.dump(yaml_data, default_flow_style=False)
+                yaml_string = yaml_string.replace('\'\'', '""')
+                yaml_file.write(yaml_string)
 
         return module.exit_json(changed=(len(changes) > 0), changes=changes)
 
     # ignore broad-except error to avoid stack trace to ansible user
     # pylint: disable=broad-except
-    except Exception, e:
-        return module.fail_json(msg=str(e))
+    except Exception as error:
+        return module.fail_json(msg=str(error))
 
-# ignore pylint errors related to the module_utils import
-# pylint: disable=redefined-builtin, unused-wildcard-import, wildcard-import
-# import module snippets
-from ansible.module_utils.basic import *
 
 if __name__ == '__main__':
     main()