Ensure correct serialize is called for List and Map attributes (#286)

Author: Jean-Charles BERTIN

pynamodb/attributes.py

@@ -487,10 +487,13 @@ def serialize(self, values):
         rval = {}
         for k in values:
             v = values[k]
-            attr_class = _get_class_for_serialize(v)
-            attr_key = _get_key_for_serialize(v)
+            attr_class = self._get_serialize_class(k, v)
             if attr_class is None:
                 continue
+            if attr_class.attr_type:
+                attr_key = ATTR_TYPE_MAP[attr_class.attr_type]
+            else:
+                attr_key = _get_key_for_serialize(v)
 
             # If this is a subclassed MapAttribute, there may be an alternate attr name
             attr = self._get_attributes().get(k)
@@ -531,6 +534,12 @@ def as_dict(self):
             result[key] = value.as_dict() if isinstance(value, MapAttribute) else value
         return result
 
+    @classmethod
+    def _get_serialize_class(cls, key, value):
+        if not cls.is_raw():
+            return cls._get_attributes().get(key)
+        return _get_class_for_serialize(value)
+
     @classmethod
     def _get_deserialize_class(cls, key, value):
         if not cls.is_raw():
@@ -595,8 +604,13 @@ def serialize(self, values):
         """
         rval = []
         for v in values:
-            attr_class = _get_class_for_serialize(v)
-            attr_key = _get_key_for_serialize(v)
+            attr_class = (self.element_type()
+                          if self.element_type
+                          else _get_class_for_serialize(v))
+            if attr_class.attr_type:
+                attr_key = ATTR_TYPE_MAP[attr_class.attr_type]
+            else:
+                attr_key = _get_key_for_serialize(v)
             rval.append({attr_key: attr_class.serialize(v)})
         return rval
 

pynamodb/tests/test_attributes.py

@@ -10,7 +10,7 @@
 from dateutil.parser import parse
 from dateutil.tz import tzutc
 
-from mock import patch
+from mock import patch, Mock, call
 
 from pynamodb.compat import CompatTestCase as TestCase
 from pynamodb.constants import UTC, DATETIME_FORMAT
@@ -19,7 +19,7 @@
 from pynamodb.attributes import (
     BinarySetAttribute, BinaryAttribute, NumberSetAttribute, NumberAttribute,
     UnicodeAttribute, UnicodeSetAttribute, UTCDateTimeAttribute, BooleanAttribute, LegacyBooleanAttribute,
-    MapAttribute, ListAttribute,
+    MapAttribute, ListAttribute, Attribute,
     JSONAttribute, DEFAULT_ENCODING, NUMBER, STRING, STRING_SET, NUMBER_SET, BINARY_SET,
     BINARY, MAP, LIST, BOOLEAN, _get_value_for_deserialize)
 
@@ -395,7 +395,7 @@ def test_unicode_set_deserialize(self):
             value
         )
 
-    def test_unicode_set_deserialize(self):
+    def test_unicode_set_deserialize_old_way(self):
         """
         UnicodeSetAttribute.deserialize old way
         """
@@ -655,6 +655,34 @@ class SomeModel(Model):
         self.assertEqual(json.dumps({'map_attr': {'foo': 'bar'}}),
                          json.dumps(item.typed_map.as_dict()))
 
+    def test_json_serialize(self):
+        class JSONMapAttribute(MapAttribute):
+            arbitrary_data = JSONAttribute()
+
+            def __eq__(self, other):
+                return self.arbitrary_data == other.arbitrary_data
+
+        item = {'foo': 'bar', 'bool': True, 'number': 3.141}
+        json_map = JSONMapAttribute(arbitrary_data=item)
+        serialized = json_map.serialize(json_map)
+        deserialized = json_map.deserialize(serialized)
+        self.assertTrue(isinstance(deserialized, JSONMapAttribute))
+        self.assertEqual(deserialized, json_map)
+        self.assertEqual(deserialized.arbitrary_data, item)
+
+    def test_serialize_datetime(self):
+        class CustomMapAttribute(MapAttribute):
+            date_attr = UTCDateTimeAttribute()
+
+        cm = CustomMapAttribute(date_attr=datetime(2017, 1, 1))
+        serialized_datetime = cm.serialize(cm)
+        expected_serialized_value = {
+            'date_attr': {
+                'S': u'2017-01-01T00:00:00.000000+0000'
+            }
+        }
+        self.assertEquals(serialized_datetime, expected_serialized_value)
+
 
 class ValueDeserializeTestCase(TestCase):
     def test__get_value_for_deserialize(self):
@@ -728,8 +756,8 @@ def __lt__(self, other):
                 return self.name < other.name
 
             def __eq__(self, other):
-                return self.name == other.name and \
-                       self.age == other.age
+                return (self.name == other.name and
+                        self.age == other.age)
 
         person1 = Person()
         person1.name = 'john'
@@ -738,13 +766,65 @@ def __eq__(self, other):
         person2 = Person()
         person2.name = 'Dana'
         person2.age = 41
+
         inp = [person1, person2]
 
         list_attribute = ListAttribute(default=[], of=Person)
         serialized = list_attribute.serialize(inp)
         deserialized = list_attribute.deserialize(serialized)
         self.assertEqual(sorted(deserialized), sorted(inp))
 
+    def test_list_of_map_with_of_and_custom_attribute(self):
+
+        # Create a couple of mock functions to use
+        # to test that the CustomAttribute serialize/deserialize are called
+        serialize_mock = Mock()
+        deserialize_mock = Mock()
+
+        class CustomAttribute(Attribute):
+            attr_type = STRING
+
+            def serialize(self, value):
+                serialize_mock(value)
+                return value.upper()
+
+            def deserialize(self, value):
+                deserialize_mock(value)
+                return value.lower()
+
+        class CustomMapAttribute(MapAttribute):
+            custom = CustomAttribute()
+
+            def __lt__(self, other):
+                return self.custom < other.custom
+
+            def __eq__(self, other):
+                return self.custom == other.custom
+
+        attribute1 = CustomMapAttribute()
+        attribute1.custom = 'test-value1'
+
+        attribute2 = CustomMapAttribute()
+        attribute2.custom = 'test-value2'
+
+        inp = [attribute1, attribute2]
+
+        list_attribute = ListAttribute(default=[], of=CustomMapAttribute)
+        serialized = list_attribute.serialize(inp)
+        deserialized = list_attribute.deserialize(serialized)
+        self.assertEqual(sorted(deserialized), sorted(inp))
+
+        # Confirm that the the serialize/deserialize are called
+        # with the expected values
+        serialize_mock.assert_has_calls([
+            call('test-value1'),
+            call('test-value2'),
+        ])
+        deserialize_mock.assert_has_calls([
+            call('TEST-VALUE1'),
+            call('TEST-VALUE2'),
+        ])
+
     def test_list_of_unicode_with_of(self):
         with self.assertRaises(ValueError):
             ListAttribute(default=[], of=UnicodeAttribute)

requirements-dev.txt

@@ -1,5 +1,6 @@
 -rrequirements.txt
 coverage==3.7.1
+mock==2.0.0
 pytest==3.0.7
 pytest-cov==2.4.0
 python-coveralls==2.5.0