Extract projection_expression adjustment

moto/dynamodb2/responses.py

@@ -308,20 +308,9 @@ class DynamoHandler(BaseResponse):
         projection_expression = self.body.get('ProjectionExpression')
         expression_attribute_names = self.body.get('ExpressionAttributeNames', {})
 
-        if projection_expression and expression_attribute_names:
+        projection_expression = self._adjust_projection_expression(
-            expressions = [x.strip() for x in projection_expression.split(',')]
+            projection_expression, expression_attribute_names
-            projection_expression = None
+        )
-            for expression in expressions:
-                if projection_expression is not None:
-                    projection_expression = projection_expression + ", "
-                else:
-                    projection_expression = ""
-
-                if expression in expression_attribute_names:
-                    projection_expression = projection_expression + \
-                        expression_attribute_names[expression]
-                else:
-                    projection_expression = projection_expression + expression
 
         try:
             item = self.dynamodb_backend.get_item(name, key, projection_expression)
@@ -359,20 +348,9 @@ class DynamoHandler(BaseResponse):
             projection_expression = table_request.get('ProjectionExpression')
             expression_attribute_names = table_request.get('ExpressionAttributeNames', {})
 
-            if projection_expression and expression_attribute_names:
+            projection_expression = self._adjust_projection_expression(
-                expressions = [x.strip() for x in projection_expression.split(',')]
+                projection_expression, expression_attribute_names
-                projection_expression = None
+            )
-                for expression in expressions:
-                    if projection_expression is not None:
-                        projection_expression = projection_expression + ", "
-                    else:
-                        projection_expression = ""
-
-                    if expression in expression_attribute_names:
-                        projection_expression = projection_expression + \
-                            expression_attribute_names[expression]
-                    else:
-                        projection_expression = projection_expression + expression
 
             results["Responses"][table_name] = []
             for key in keys:
@@ -406,20 +384,9 @@ class DynamoHandler(BaseResponse):
         filter_expression = self.body.get('FilterExpression')
         expression_attribute_values = self.body.get('ExpressionAttributeValues', {})
 
-        if projection_expression and expression_attribute_names:
+        projection_expression = self._adjust_projection_expression(
-            expressions = [x.strip() for x in projection_expression.split(',')]
+            projection_expression, expression_attribute_names
-            projection_expression = None
+        )
-            for expression in expressions:
-                if projection_expression is not None:
-                    projection_expression = projection_expression + ", "
-                else:
-                    projection_expression = ""
-
-                if expression in expression_attribute_names:
-                    projection_expression = projection_expression + \
-                        expression_attribute_names[expression]
-                else:
-                    projection_expression = projection_expression + expression
 
         filter_kwargs = {}
 
@@ -555,6 +522,25 @@ class DynamoHandler(BaseResponse):
 
         return dynamo_json_dump(result)
 
+    def _adjust_projection_expression(self, projection_expression, expression_attribute_names):
+        if projection_expression and expression_attribute_names:
+            expressions = [x.strip() for x in projection_expression.split(',')]
+            projection_expr = None
+            for expression in expressions:
+                if projection_expr is not None:
+                    projection_expr = projection_expr + ", "
+                else:
+                    projection_expr = ""
+
+                if expression in expression_attribute_names:
+                    projection_expr = projection_expr + \
+                        expression_attribute_names[expression]
+                else:
+                    projection_expr = projection_expr + expression
+            return projection_expr
+
+        return projection_expression
+
     def scan(self):
         name = self.body['TableName']
 

tests/test_dynamodb2/test_dynamodb.py

@@ -421,10 +421,10 @@ def test_basic_projection_expression_using_get_item():
         ProjectionExpression='body, subject'
     )
 
-    assert result['Item'] == {
+    result['Item'].should.be.equal({
         'subject': '123',
         'body': 'some test message'
-    }
+    })
 
     # The projection expression should not remove data from storage
     result = table.get_item(
@@ -434,11 +434,11 @@ def test_basic_projection_expression_using_get_item():
         }
     )
 
-    assert result['Item'] == {
+    result['Item'].should.be.equal({
         'forum_name': 'the-key',
         'subject': '123',
         'body': 'some test message'
-    }
+    })
 
 
 @mock_dynamodb2
@@ -670,11 +670,11 @@ def test_basic_projection_expression_using_get_item_with_attr_expression_names()
             },
     )
 
-    assert result['Item'] == {
+    result['Item'].should.be.equal({
         'subject': '123',
         'body': 'some test message',
         'attachment': 'something'
-    }
+    })
 
 
 @mock_dynamodb2
@@ -2392,9 +2392,10 @@ def test_batch_items_with_basic_projection_expression():
             'ProjectionExpression': 'username'
         }
     })['Responses']['users']
-    assert len(returned_items) == 3
+
-    assert [item['username']['S'] for item in returned_items] == ['user1', 'user2', 'user3']
+    returned_items.should.have.length_of(3)
-    assert [item.get('foo') for item in returned_items] == [None, None, None]
+    [item['username']['S'] for item in returned_items].should.be.equal(['user1', 'user2', 'user3'])
+    [item.get('foo') for item in returned_items].should.be.equal([None, None, None])
 
     # The projection expression should not remove data from storage
     returned_items = dynamodb.batch_get_item(RequestItems = {
@@ -2411,8 +2412,9 @@ def test_batch_items_with_basic_projection_expression():
             'ConsistentRead': True
         }
     })['Responses']['users']
-    assert [item['username']['S'] for item in returned_items] == ['user1', 'user2', 'user3']
+
-    assert [item['foo']['S'] for item in returned_items] == ['bar', 'bar', 'bar']
+    [item['username']['S'] for item in returned_items].should.be.equal(['user1', 'user2', 'user3'])
+    [item['foo']['S'] for item in returned_items].should.be.equal(['bar', 'bar', 'bar'])
 
 
 @mock_dynamodb2
@@ -2436,9 +2438,10 @@ def test_batch_items_with_basic_projection_expression_and_attr_expression_names(
             },
         }
     })['Responses']['users']
-    assert len(returned_items) == 3
+
-    assert [item['username']['S'] for item in returned_items] == ['user1', 'user2', 'user3']
+    returned_items.should.have.length_of(3)
-    assert [item.get('foo') for item in returned_items] == [None, None, None]
+    [item['username']['S'] for item in returned_items].should.be.equal(['user1', 'user2', 'user3'])
+    [item.get('foo') for item in returned_items].should.be.equal([None, None, None])
 
 
 @mock_dynamodb2