Nicer update syntax for adding and removing subsets. Fixes pynamodb#389 (pynamodb#408)

jpinner-lyft · web-flow · commit 1828bda52376 · 2017-11-09T09:37:45.000-08:00
diff --git a/pynamodb/attributes.py b/pynamodb/attributes.py
@@ -163,11 +163,11 @@ def set(self, value):
     def remove(self):
         return Path(self).remove()
 
-    def add(self, value):
-        return Path(self).add(value)
+    def add(self, *values):
+        return Path(self).add(*values)
 
-    def delete(self, value):
-        return Path(self).delete(value)
+    def delete(self, *values):
+        return Path(self).delete(*values)
 
 
 class AttributeContainerMeta(type):
diff --git a/pynamodb/expressions/operand.py b/pynamodb/expressions/operand.py
@@ -275,12 +275,14 @@ def remove(self):
         # Returns an update action that removes this attribute from the item
         return RemoveAction(self)
 
-    def add(self, value):
-        # Returns an update action that appends the given value to a set or mathematically adds it to a number
+    def add(self, *values):
+        # Returns an update action that appends the given values to a set or mathematically adds a value to a number
+        value = values[0] if len(values) == 1 else values
         return AddAction(self, self._to_operand(value))
 
-    def delete(self, value):
-        # Returns an update action that removes the given value from a set attribute
+    def delete(self, *values):
+        # Returns an update action that removes the given values from a set attribute
+        value = values[0] if len(values) == 1 else values
         return DeleteAction(self, self._to_operand(value))
 
     def exists(self):
diff --git a/pynamodb/tests/test_expressions.py b/pynamodb/tests/test_expressions.py
@@ -475,6 +475,14 @@ def test_add_action(self):
         assert expression_attribute_values == {':0': {'N': '0'}}
 
     def test_add_action_set(self):
+        action = NumberSetAttribute(attr_name='foo').add(0, 1)
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = action.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0 :0"
+        assert placeholder_names == {'foo': '#0'}
+        assert expression_attribute_values == {':0': {'NS': ['0', '1']}}
+
+    def test_add_action_serialized(self):
         action = NumberSetAttribute(attr_name='foo').add({'NS': ['0']})
         placeholder_names, expression_attribute_values = {}, {}
         expression = action.serialize(placeholder_names, expression_attribute_values)
@@ -487,6 +495,22 @@ def test_add_action_list(self):
             Path('foo').add({'L': [{'N': '0'}]})
 
     def test_delete_action(self):
+        action = NumberSetAttribute(attr_name='foo').delete(0, 1)
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = action.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0 :0"
+        assert placeholder_names == {'foo': '#0'}
+        assert expression_attribute_values == {':0': {'NS': ['0', '1']}}
+
+    def test_delete_action_set(self):
+        action = NumberSetAttribute(attr_name='foo').delete(set([0, 1]))
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = action.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0 :0"
+        assert placeholder_names == {'foo': '#0'}
+        assert expression_attribute_values == {':0': {'NS': ['0', '1']}}
+
+    def test_delete_action_serialized(self):
         action = NumberSetAttribute(attr_name='foo').delete({'NS': ['0']})
         placeholder_names, expression_attribute_values = {}, {}
         expression = action.serialize(placeholder_names, expression_attribute_values)
diff --git a/pynamodb/tests/test_model.py b/pynamodb/tests/test_model.py
@@ -966,7 +966,8 @@ def test_update(self):
                 SimpleUserModel.views.remove(),
                 SimpleUserModel.is_active.set(None),
                 SimpleUserModel.signature.set(None),
-                SimpleUserModel.custom_aliases.set(['bob'])
+                SimpleUserModel.custom_aliases.set(['bob']),
+                SimpleUserModel.numbers.delete(0, 1)
             ])
 
             args = req.call_args[0][1]
@@ -978,13 +979,14 @@ def test_update(self):
                         'S': 'foo'
                     }
                 },
-                'UpdateExpression': 'SET #0 = :0, #1 = :1, #2 = :2, #3 = :3 REMOVE #4',
+                'UpdateExpression': 'SET #0 = :0, #1 = :1, #2 = :2, #3 = :3 REMOVE #4 DELETE #5 :4',
                 'ExpressionAttributeNames': {
                     '#0': 'email',
                     '#1': 'is_active',
                     '#2': 'signature',
                     '#3': 'aliases',
-                    '#4': 'views'
+                    '#4': 'views',
+                    '#5': 'numbers'
                 },
                 'ExpressionAttributeValues': {
                     ':0': {
@@ -997,7 +999,10 @@ def test_update(self):
                         'NULL': True
                     },
                     ':3': {
-                        'SS': set(['bob'])
+                        'SS': ['bob']
+                    },
+                    ':4': {
+                        'NS': ['0', '1']
                     }
                 },
                 'ReturnConsumedCapacity': 'TOTAL'
