add support for nested keys

salt/manager/tools/sbin/so-yaml.py

@@ -16,12 +16,12 @@ lockFile = "/tmp/so-yaml.lock"
 def showUsage(args):
     print('Usage: {} <COMMAND> <YAML_FILE> [ARGS...]'.format(sys.argv[0]))
     print('  General commands:')
-    print('    remove         - Removes a yaml top-level key, if it exists. Requires KEY arg.')
+    print('    remove         - Removes a yaml key, if it exists. Requires KEY arg.')
     print('    help           - Prints this usage information.')
     print('')
     print('  Where:')
     print('   YAML_FILE       - Path to the file that will be modified. Ex: /opt/so/conf/service/conf.yaml')
-    print('   KEY             - Top level key only, does not support dot-notations for nested keys at this time. Ex: level1')
+    print('   KEY             - YAML key, does not support \' or " characters at this time. Ex: level1.level2')
     sys.exit(1)
 
 
@@ -36,6 +36,13 @@ def writeYaml(filename, content):
     return yaml.dump(content, file)
 
 
+def removeKey(content, key):
+    pieces = key.split(".", 1)
+    if len(pieces) > 1:
+        removeKey(content[pieces[0]], pieces[1])
+    else:
+        content.pop(key, None)
+
 def remove(args):
     if len(args) != 2:
         print('Missing filename or key arg', file=sys.stderr)
@@ -43,11 +50,12 @@ def remove(args):
         return
 
     filename = args[0]
+    key = args[1]
+
     content = loadYaml(filename)
-
+    removeKey(content, key)
-    content.pop(args[1], None)
-
     writeYaml(filename, content)
+
     return 0
 
 

salt/manager/tools/sbin/so-yaml_test.py

@@ -57,6 +57,36 @@ class TestRemove(unittest.TestCase):
         expected = "key2: false\n"
         self.assertEqual(actual, expected)
 
+    def test_remove_nested(self):
+        filename = "/tmp/so-yaml_test-remove.yaml"
+        file = open(filename, "w")
+        file.write("{key1: { child1: 123, child2: abc }, key2: false}")
+        file.close()
+
+        soyaml.remove([filename, "key1.child2"])
+
+        file = open(filename, "r")
+        actual = file.read()
+        file.close()
+
+        expected = "key1:\n  child1: 123\nkey2: false\n"
+        self.assertEqual(actual, expected)
+
+    def test_remove_nested_deep(self):
+        filename = "/tmp/so-yaml_test-remove.yaml"
+        file = open(filename, "w")
+        file.write("{key1: { child1: 123, child2: { deep1: 45, deep2: ab } }, key2: false}")
+        file.close()
+
+        soyaml.remove([filename, "key1.child2.deep1"])
+
+        file = open(filename, "r")
+        actual = file.read()
+        file.close()
+
+        expected = "key1:\n  child1: 123\n  child2:\n    deep2: ab\nkey2: false\n"
+        self.assertEqual(actual, expected)
+
     def test_remove_missing_args(self):
         with patch('sys.exit', new=MagicMock()) as sysmock:
             with patch('sys.stderr', new=StringIO()) as mock_stdout: