Enable polymorphism by serializing all attributes in subclasses. (#842)

Author: jpinner-lyft

pynamodb/attributes.py

@@ -313,6 +313,24 @@ def _set_attributes(self, **attributes: Attribute) -> None:
                 raise ValueError("Attribute {} specified does not exist".format(attr_name))
             setattr(self, attr_name, attr_value)
 
+    def _serialize(self, null_check=True) -> Dict[str, Dict[str, Any]]:
+        """
+        Serialize attribute values for DynamoDB
+        """
+        attribute_values: Dict[str, Dict[str, Any]] = {}
+        for name, attr in self.get_attributes().items():
+            value = getattr(self, name)
+            if isinstance(value, MapAttribute) and not value.validate():
+                raise ValueError("Attribute '{}' is not correctly typed".format(name))
+
+            attr_value = attr.serialize(value) if value is not None else None
+            if null_check and attr_value is None and not attr.null:
+                raise ValueError("Attribute '{}' cannot be None".format(name))
+
+            if attr_value is not None:
+                attribute_values[attr.attr_name] = {attr.attr_type: attr_value}
+        return attribute_values
+
     def _deserialize(self, attribute_values: Dict[str, Dict[str, Any]]) -> None:
         """
         Sets attributes sent back from DynamoDB on this object
@@ -807,25 +825,34 @@ def validate(self):
         return all(self.is_correctly_typed(k, v) for k, v in self.get_attributes().items())
 
     def serialize(self, values):
-        rval = {}
-        for k in values:
-            v = values[k]
-            if self._should_skip(v):
-                continue
-            attr_class = self._get_serialize_class(k, v)
-            if attr_class is None:
-                continue
-
-            # If this is a subclassed MapAttribute, there may be an alternate attr name
-            attr_name = attr_class.attr_name if not self.is_raw() else k
+        if not self.is_raw():
+            # This is a subclassed MapAttribute that acts as an AttributeContainer.
+            # Serialize the values based on the attributes in the class.
 
-            serialized = attr_class.serialize(v)
-            if self._should_skip(serialized):
-                # Check after we serialize in case the serialized value is null
-                continue
+            if not isinstance(values, type(self)):
+                # Copy the values onto an instance of the class for serialization.
+                instance = type(self)()
+                instance.attribute_values = {}  # clear any defaults
+                for name in values:
+                    if name in self.get_attributes():
+                        setattr(instance, name, values[name])
+                values = instance
 
-            rval[attr_name] = {attr_class.attr_type: serialized}
+            return values._serialize()
 
+        # Continue to serialize NULL values in "raw" map attributes for backwards compatibility.
+        # This special case behavior for "raw" attributes should be removed in the future.
+        rval = {}
+        for attr_name in values:
+            v = values[attr_name]
+            attr_class = _get_class_for_serialize(v)
+            attr_type = attr_class.attr_type
+            attr_value = attr_class.serialize(v)
+            if attr_value is None:
+                # When attribute values serialize to "None" (e.g. empty sets) we store {"NULL": True} in DynamoDB.
+                attr_type = NULL
+                attr_value = True
+            rval[attr_name] = {attr_type: attr_value}
         return rval
 
     def deserialize(self, values):
@@ -853,17 +880,6 @@ def as_dict(self):
             result[key] = value.as_dict() if isinstance(value, MapAttribute) else value
         return result
 
-    def _should_skip(self, value):
-        # Continue to serialize NULL values in "raw" map attributes for backwards compatibility.
-        # This special case behavior for "raw" attributes should be removed in the future.
-        return not self.is_raw() and value is None
-
-    @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)
-
 
 def _get_class_for_serialize(value):
     if value is None:

pynamodb/models.py

@@ -995,15 +995,15 @@ def _handle_version_attribute(self, serialized_attributes, actions=None):
                 actions.append(version_attribute.add(1))
             elif snake_to_camel_case(ATTRIBUTES) in serialized_attributes:
                 serialized_attributes[snake_to_camel_case(ATTRIBUTES)][version_attribute.attr_name] = self._serialize_value(
-                    version_attribute, version_attribute_value + 1, null_check=True
+                    version_attribute, version_attribute_value + 1
                 )
         else:
             version_condition = version_attribute.does_not_exist()
             if actions:
                 actions.append(version_attribute.set(1))
             elif snake_to_camel_case(ATTRIBUTES) in serialized_attributes:
                 serialized_attributes[snake_to_camel_case(ATTRIBUTES)][version_attribute.attr_name] = self._serialize_value(
-                    version_attribute, 1, null_check=True
+                    version_attribute, 1
                 )
 
         return version_condition
@@ -1107,53 +1107,40 @@ def _get_connection(cls) -> TableConnection:
                                               aws_session_token=cls.Meta.aws_session_token)
         return cls._connection
 
-    def _serialize(self, attr_map=False, null_check=True) -> Dict[str, Any]:
+    def _serialize(self, null_check=True, attr_map=False) -> Dict[str, Dict[str, Any]]:
         """
         Serializes all model attributes for use with DynamoDB
 
-        :param attr_map: If True, then attributes are returned
         :param null_check: If True, then attributes are checked for null
+        :param attr_map: If True, then attributes are returned
         """
         attributes = snake_to_camel_case(ATTRIBUTES)
-        attrs: Dict[str, Dict] = {attributes: {}}
-        for name, attr in self.get_attributes().items():
-            value = getattr(self, name)
-            if isinstance(value, MapAttribute):
-                if not value.validate():
-                    raise ValueError("Attribute '{}' is not correctly typed".format(attr.attr_name))
-
-            serialized = self._serialize_value(attr, value, null_check)
-            if NULL in serialized:
-                continue
-
-            if attr_map:
-                attrs[attributes][attr.attr_name] = serialized
-            else:
-                if attr.is_hash_key:
-                    attrs[HASH] = serialized[attr.attr_type]
-                elif attr.is_range_key:
-                    attrs[RANGE] = serialized[attr.attr_type]
-                else:
-                    attrs[attributes][attr.attr_name] = serialized
-
+        attrs: Dict[str, Dict] = {attributes: super()._serialize(null_check)}
+        if not attr_map:
+            hash_key_attribute = self._hash_key_attribute()
+            hash_key_attribute_value = attrs[attributes].pop(hash_key_attribute.attr_name, None)
+            if hash_key_attribute_value is not None:
+                attrs[HASH] = hash_key_attribute_value[hash_key_attribute.attr_type]
+            range_key_attribute = self._range_key_attribute()
+            if range_key_attribute:
+                range_key_attribute_value = attrs[attributes].pop(range_key_attribute.attr_name, None)
+                if range_key_attribute_value is not None:
+                    attrs[RANGE] = range_key_attribute_value[range_key_attribute.attr_type]
         return attrs
 
     @classmethod
-    def _serialize_value(cls, attr, value, null_check=True):
+    def _serialize_value(cls, attr, value):
         """
         Serializes a value for use with DynamoDB
 
         :param attr: an instance of `Attribute` for serialization
         :param value: a value to be serialized
         :param null_check: If True, then attributes are checked for null
         """
-        if value is None:
-            serialized = None
-        else:
-            serialized = attr.serialize(value)
+        serialized = attr.serialize(value)
 
         if serialized is None:
-            if not attr.null and null_check:
+            if not attr.null:
                 raise ValueError("Attribute '{}' cannot be None".format(attr.attr_name))
             return {NULL: True}
 

tests/test_attributes.py

@@ -614,19 +614,21 @@ def test_null_attribute_raw_map(self):
 
     def test_null_attribute_subclassed_map(self):
         null_attribute = {
-            'map_field': None
+            'map_field': {},
+            'string_set_field': None
         }
         attr = DefaultsMap()
         serialized = attr.serialize(null_attribute)
-        assert serialized == {}
+        assert serialized == {'map_field': {'M': {}}}
 
     def test_null_attribute_map_after_serialization(self):
         null_attribute = {
+            'map_field': {},
             'string_set_field': {},
         }
         attr = DefaultsMap()
         serialized = attr.serialize(null_attribute)
-        assert serialized == {}
+        assert serialized == {'map_field': {'M': {}}}
 
     def test_map_of_map(self):
         attribute = {
@@ -889,6 +891,29 @@ class MyModel(Model):
         assert mid_map_b_map_attr.attr_name == 'dyn_map_attr'
         assert mid_map_b_map_attr.attr_path == ['dyn_out_map', 'mid_map_b', 'dyn_in_map_b', 'dyn_map_attr']
 
+    def test_required_elements(self):
+        class InnerMapAttribute(MapAttribute):
+            foo = UnicodeAttribute()
+
+        class OuterMapAttribute(MapAttribute):
+            inner_map = InnerMapAttribute()
+
+        outer_map_attribute = OuterMapAttribute()
+        with pytest.raises(ValueError):
+            outer_map_attribute.serialize(outer_map_attribute)
+
+        outer_map_attribute = OuterMapAttribute(inner_map={})
+        with pytest.raises(ValueError):
+            outer_map_attribute.serialize(outer_map_attribute)
+
+        outer_map_attribute = OuterMapAttribute(inner_map=MapAttribute())
+        with pytest.raises(ValueError):
+            outer_map_attribute.serialize(outer_map_attribute)
+
+        outer_map_attribute = OuterMapAttribute(inner_map={'foo': 'bar'})
+        serialized = outer_map_attribute.serialize(outer_map_attribute)
+        assert serialized == {'inner_map': {'M': {'foo': {'S': 'bar'}}}}
+
 
 class TestListAttribute: