From ab375b212696b867dbbbc24b31275f2fbfea3409 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Wed, 9 Nov 2016 12:50:30 +0100
Subject: [PATCH] Moved node_setattr() from Attract to Pillar

---
 pillar/api/utils/__init__.py | 21 +++++++++++++
 tests/test_api/test_utils.py | 58 ++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/pillar/api/utils/__init__.py b/pillar/api/utils/__init__.py
index c096cc08..3af84d31 100644
--- a/pillar/api/utils/__init__.py
+++ b/pillar/api/utils/__init__.py
@@ -16,6 +16,27 @@ import pymongo.results
 log = logging.getLogger(__name__)
 
 
+def node_setattr(node, key, value):
+    """Sets a node property by dotted key.
+
+    Modifies the node in-place. Deletes None values.
+
+    :type node: dict
+    :type key: str
+    :param value: the value to set, or None to delete the key.
+    """
+
+    set_on = node
+    while key and '.' in key:
+        head, key = key.split('.', 1)
+        set_on = set_on[head]
+
+    if value is None:
+        set_on.pop(key, None)
+    else:
+        set_on[key] = value
+
+
 def remove_private_keys(document):
     """Removes any key that starts with an underscore, returns result as new
     dictionary.
diff --git a/tests/test_api/test_utils.py b/tests/test_api/test_utils.py
index 3b8d56d0..70860bf2 100644
--- a/tests/test_api/test_utils.py
+++ b/tests/test_api/test_utils.py
@@ -105,3 +105,61 @@ class DocDiffTest(unittest.TestCase):
         self.assertEqual({('props.status1', u'todo', DoesNotExist),
                           ('props.status2', DoesNotExist, u'todo')},
                          set(diff))
+
+
+class NodeSetattrTest(unittest.TestCase):
+    def test_simple(self):
+        from pillar.api.utils import node_setattr
+
+        node = {}
+        node_setattr(node, 'a', 5)
+        self.assertEqual({'a': 5}, node)
+
+        node_setattr(node, 'b', {'complexer': 'value'})
+        self.assertEqual({'a': 5, 'b': {'complexer': 'value'}}, node)
+
+    def test_dotted(self):
+        from pillar.api.utils import node_setattr
+
+        node = {}
+        self.assertRaises(KeyError, node_setattr, node, 'a.b', 5)
+
+        node = {'b': {}}
+        node_setattr(node, 'b.simple', 'value')
+        self.assertEqual({'b': {'simple': 'value'}}, node)
+
+        node_setattr(node, 'b.complex', {'yes': 'value'})
+        self.assertEqual({'b': {'simple': 'value',
+                                'complex': {'yes': 'value'}}}, node)
+
+        node_setattr(node, 'b.complex', {'yes': 5})
+        self.assertEqual({'b': {'simple': 'value',
+                                'complex': {'yes': 5}}}, node)
+
+    def test_none_simple(self):
+        from pillar.api.utils import node_setattr
+
+        node = {}
+        node_setattr(node, 'a', None)
+        node_setattr(node, None, 'b')
+        self.assertEqual({None: 'b'}, node)
+
+    def test_none_dotted(self):
+        from pillar.api.utils import node_setattr
+
+        node = {}
+        self.assertRaises(KeyError, node_setattr, node, 'a.b', None)
+
+        node = {'b': {}}
+        node_setattr(node, 'b.simple', None)
+        self.assertEqual({'b': {}}, node)
+
+        node_setattr(node, 'b.complex', {'yes': None})
+        self.assertEqual({'b': {'complex': {'yes': None}}}, node)
+
+        node_setattr(node, 'b.complex.yes', None)
+        self.assertEqual({'b': {'complex': {}}}, node)
+
+        node_setattr(node, 'b.complex', {None: 5})
+        self.assertEqual({'b': {'complex': {None: 5}}}, node)
+