diff --git a/boto3_helpers/dynamodb.py b/boto3_helpers/dynamodb.py index b1063c0..ff1ecc7 100644 --- a/boto3_helpers/dynamodb.py +++ b/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()} diff --git a/tests/test_dynamodb.py b/tests/test_dynamodb.py index 1a58204..2d64ab7 100644 --- a/tests/test_dynamodb.py +++ b/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)