Merge pull request #18 from openwurl/ddb-marshal

Add dynamodb.fix_numbers

boto3_helpers/dynamodb.py

@@ -1,7 +1,22 @@
 from boto3 import resource as boto3_resource
+from boto3.dynamodb.types import TypeDeserializer, TypeSerializer
+
 from time import sleep
 
 
+class CustomTypeDeserializer(TypeDeserializer):
+    def __init__(self, *args, use_decimal=False, **kwargs):
+        self.use_decimal = use_decimal
+        super().__init__(*args, **kwargs)
+
+    def _deserialize_n(self, value):
+        if self.use_decimal:
+            return super()._deserialize_n(value)
+
+        ret = float(value)
+        return int(ret) if ret.is_integer() else ret
+
+
 def _table_or_name(x):
     if isinstance(x, str):
         return boto3_resource('dynamodb').Table(x)
@@ -201,3 +216,28 @@ def batch_yield_items(
             break
         sleep(min(backoff_base * (2**i), backoff_max))
         i += 1
+
+
+def fix_numbers(item):
+    """``boto3`` DB infamously deserializes numeric types from DynamoDB to
+    Python ``Decimal`` objects. This function changes these objects into
+    ``int`` objects and ``float`` objects.
+
+    .. code-block:: python
+
+        from boto3 import resource as boto3_resource
+        from boto3_helpers.dynamodb import fix_numbers
+
+        ddb_resource = boto3_resource('dynamodb')
+        ddb_table = ddb_resource.Table('example-table')
+        resp = ddb_table.get_item(Key={'primary_key': 'FirstKey'})
+        item = resp['Item']
+        fixed_item = fix_numbers(item)
+
+    Note that ``float`` objects may not be appropriate for all numeric computing needs,
+    so think about what your application needs before using this function.
+    """
+    s = TypeSerializer().serialize
+    d = CustomTypeDeserializer().deserialize
+    wire_format = {k: s(v) for k, v in item.items()}
+    return {k: d(v) for k, v in wire_format.items()}

tests/test_dynamodb.py

@@ -8,6 +8,7 @@
 
 from boto3_helpers.dynamodb import (
     batch_yield_items,
+    fix_numbers,
     query_table,
     scan_table,
     update_attributes,
@@ -234,3 +235,44 @@ def test_batch_yield_items_batch_size(self, mock_boto3_resource, mock_sleep):
         )
         mock_boto3_resource.assert_called_once_with('dynamodb')
         self.assertEqual(mock_boto3_resource.return_value.batch_get_item.call_count, 4)
+
+    def test_fix_numbers(self):
+        # Set up the stubber
+        ddb_resource = boto3_resource('dynamodb', region_name='not-a-region')
+        ddb_table = ddb_resource.Table('test-table')
+        stubber = Stubber(ddb_resource.meta.client)
+
+        # The first query will return a single page of results
+        resp = {
+            'Item': {
+                'string_set': {'SS': ['ss_1', 'ss_2']},
+                'number_int': {'N': '1'},
+                'number_set': {'NS': ['1.1', '1']},
+                'TestKey': {'S': 'test-key'},
+                'list_value': {'L': [{'S': 'sl_1'}, {'N': '1'}]},
+                'bool_value': {'BOOL': True},
+                'null_value': {'NULL': True},
+                'number_float': {'N': '1.1'},
+                'map_value': {'M': {'n_key': {'N': '1.1'}, 's_key': {'S': 's_value'}}},
+            }
+        }
+        params = {'TableName': 'test-table', 'Key': {'TestKey': 'test-key'}}
+        stubber.add_response('get_item', resp, params)
+
+        # Do the deed
+        with stubber:
+            item = ddb_table.get_item(Key={'TestKey': 'test-key'})['Item']
+            actual = fix_numbers(item)
+
+        expected = {
+            'string_set': {'ss_1', 'ss_2'},
+            'number_int': 1,
+            'number_set': {1, 1.1},
+            'TestKey': 'test-key',
+            'list_value': ['sl_1', 1],
+            'bool_value': True,
+            'null_value': None,
+            'number_float': 1.1,
+            'map_value': {'n_key': 1.1, 's_key': 's_value'},
+        }
+        self.assertEqual(actual, expected)