Skip to content

Commit

Permalink
Nicer update syntax for adding and removing subsets. Fixes #389 (#408)
Browse files Browse the repository at this point in the history
  • Loading branch information
jpinner-lyft committed Nov 9, 2017
1 parent 14e2507 commit 1828bda
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 12 deletions.
8 changes: 4 additions & 4 deletions pynamodb/attributes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
10 changes: 6 additions & 4 deletions pynamodb/expressions/operand.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
24 changes: 24 additions & 0 deletions pynamodb/tests/test_expressions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down
13 changes: 9 additions & 4 deletions pynamodb/tests/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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': {
Expand All @@ -997,7 +999,10 @@ def test_update(self):
'NULL': True
},
':3': {
'SS': set(['bob'])
'SS': ['bob']
},
':4': {
'NS': ['0', '1']
}
},
'ReturnConsumedCapacity': 'TOTAL'
Expand Down

0 comments on commit 1828bda

Please sign in to comment.