diff --git a/attract/shots/__init__.py b/attract/shots/__init__.py index 8d9fbc1..b7f8999 100644 --- a/attract/shots/__init__.py +++ b/attract/shots/__init__.py @@ -135,14 +135,18 @@ class ShotManager(object): def node_setattr(node, key, value): """Sets a node property by dotted key. - Modifies the node in-place. + Modifies the node in-place. Deletes None values. """ set_on = node while key and '.' in key: head, key = key.split('.', 1) - set_on = node[head] - set_on[key] = value + set_on = set_on[head] + + if value is None: + set_on.pop(key, None) + else: + set_on[key] = value @register_patch_handler(node_type_shot['name']) diff --git a/tests/test_shots.py b/tests/test_shots.py index 782e0ba..88f5d4d 100644 --- a/tests/test_shots.py +++ b/tests/test_shots.py @@ -157,7 +157,7 @@ class NodeSetattrTest(unittest.TestCase): node = {} node_setattr(node, 'a', None) node_setattr(node, None, 'b') - self.assertEqual({'a': None, None: 'b'}, node) + self.assertEqual({None: 'b'}, node) def test_none_dotted(self): from attract.shots import node_setattr @@ -167,15 +167,16 @@ class NodeSetattrTest(unittest.TestCase): node = {'b': {}} node_setattr(node, 'b.simple', None) - self.assertEqual({'b': {'simple': None}}, node) + self.assertEqual({'b': {}}, node) node_setattr(node, 'b.complex', {'yes': None}) - self.assertEqual({'b': {'simple': None, - 'complex': {'yes': None}}}, node) + 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': {'simple': None, - 'complex': {None: 5}}}, node) + self.assertEqual({'b': {'complex': {None: 5}}}, node) class PatchShotTest(AbstractShotTest):