Introduce DynamicMapAttribute (#868)

* Introduce DynamicMapAttribute

Co-authored-by: Justin Phillips <jphillips@lyft.com>
Co-authored-by: Garrett Heel <garrettheel@gmail.com>

docs/attributes.rst

@@ -194,3 +194,17 @@ These attributes can then be used inside of Model classes just like any other at
         model = UnicodeAttribute(null=True)
 
 `As with a model and its top-level attributes <https://github.com/pynamodb/PynamoDB/blob/master/docs/quickstart.rst#changing-items>`_, a PynamoDB MapAttribute will ignore sub-attributes it does not know about during deserialization. As a result, if the item in DynamoDB contains sub-attributes not declared as properties of the corresponding MapAttribute, save() will cause those sub-attributes to be deleted.
+
+``DynamicMapAttribute`` is a subclass of ``MapAttribute`` which allows you to mix and match defined attributes and undefined attributes.
+
+.. code-block:: python
+
+    from pynamodb.attributes import DynamicMapAttribute, UnicodeAttribute
+
+    class CarInfo(DynamicMapAttribute):
+        make = UnicodeAttribute(null=False)
+        model = UnicodeAttribute(null=True)
+
+    car = CarInfo(make='Make-A', model='Model-A', year=1975)
+    other_car = CarInfo(make='Make-A', model='Model-A', year=1975, seats=3)
+

pynamodb/attributes.py

@@ -862,20 +862,20 @@ def __getitem__(self, item: _KT) -> _VT:  # type: ignore
         # If this instance is being used as an Attribute, treat item access like the map dereference operator.
         # This provides equivalence between DynamoDB's nested attribute access for map elements (MyMap.nestedField)
         # and Python's item access for dictionaries (MyMap['nestedField']).
-        if self.is_raw():
-            return Path(self.attr_path + [str(item)])  # type: ignore
-        elif item in self._attributes:  # type: ignore
+        if item in self.get_attributes():
             return getattr(self, item)
+        elif self.is_raw():
+            return Path(self.attr_path + [str(item)])  # type: ignore
         else:
             raise AttributeError("'{}' has no attribute '{}'".format(self.__class__.__name__, item))
 
     def __setitem__(self, item, value):
         if not self._is_attribute_container():
             raise TypeError("'{}' object does not support item assignment".format(self.__class__.__name__))
-        if self.is_raw():
-            self.attribute_values[item] = value
-        elif item in self._attributes:  # type: ignore
+        if item in self.get_attributes():
             setattr(self, item, value)
+        elif self.is_raw():
+            self.attribute_values[item] = value
         else:
             raise AttributeError("'{}' has no attribute '{}'".format(self.__class__.__name__, item))
 
@@ -933,6 +933,22 @@ def validate(self, *, null_check: bool = False):
         return all(self.is_correctly_typed(k, v, null_check=null_check)
                    for k, v in self.get_attributes().items())
 
+    def _serialize_undeclared_attributes(self, values, container: Dict):
+        # 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.
+        for attr_name in values:
+            if attr_name not in self.get_attributes():
+                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
+                container[attr_name] = {attr_type: attr_value}
+        return container
+
     def serialize(self, values, *, null_check: bool = True):
         if not self.is_raw():
             # This is a subclassed MapAttribute that acts as an AttributeContainer.
@@ -949,20 +965,8 @@ def serialize(self, values, *, null_check: bool = True):
 
             return AttributeContainer._container_serialize(values, null_check=null_check)
 
-        # 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
+        # For a "raw" MapAttribute all fields are undeclared
+        return self._serialize_undeclared_attributes(values, {})
 
     def deserialize(self, values):
         """
@@ -988,6 +992,62 @@ def as_dict(self):
         return result
 
 
+class DynamicMapAttribute(MapAttribute):
+    """
+    A map attribute that supports declaring attributes (like an AttributeContainer) but will also store
+    any other values that are set on it (like a raw MapAttribute).
+
+    >>> class MyDynamicMapAttribute(DynamicMapAttribute):
+    >>>     a_date_time = UTCDateTimeAttribute()  # raw map attributes cannot serialize/deserialize datetime values
+    >>>
+    >>> dynamic_map = MyDynamicMapAttribute()
+    >>> dynamic_map.a_date_time = datetime.utcnow()
+    >>> dynamic_map.a_number = 5
+    >>> dynamic_map.serialize()  # {'a_date_time': {'S': 'xxx'}, 'a_number': {'N': '5'}}
+    """
+
+    def __setattr__(self, name, value):
+        # Set attributes via the Attribute descriptor if it exists.
+        if name in self.get_attributes():
+            object.__setattr__(self, name, value)
+        else:
+            super().__setattr__(name, value)
+
+    def serialize(self, values, *, null_check: bool = True):
+        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
+            instance._set_attributes(**values)
+            values = instance
+
+        # this serializes the class defined attributes.
+        # we do this first because we have type checks that validate the data
+        rval = AttributeContainer._container_serialize(values, null_check=null_check)
+
+        # this serializes the dynamically defined attributes
+        # we have no real type safety here so we have to dynamically construct the type to write to dynamo
+        self._serialize_undeclared_attributes(values, rval)
+
+        return rval
+
+    def deserialize(self, values):
+        # this deserializes the class defined attributes
+        # we do this first so that the we populate the defined object attributes fields properly with type safety
+        instance = self._instantiate(values)
+        # this deserializes the dynamically defined attributes
+        for attr_name, value in values.items():
+            if instance._dynamo_to_python_attr(attr_name) not in instance.get_attributes():
+                attr_type, attr_value = next(iter(value.items()))
+                instance[attr_name] = DESERIALIZE_CLASS_MAP[attr_type].deserialize(attr_value)
+        return instance
+
+    @classmethod
+    def is_raw(cls):
+        # All subclasses of DynamicMapAttribute should be treated like "raw" map attributes.
+        return True
+
+
 def _get_class_for_serialize(value):
     if value is None:
         return NullAttribute()

tests/test_attributes.py

@@ -12,7 +12,7 @@
 import pytest
 
 from pynamodb.attributes import (
-    BinarySetAttribute, BinaryAttribute, NumberSetAttribute, NumberAttribute,
+    BinarySetAttribute, BinaryAttribute, DynamicMapAttribute, NumberSetAttribute, NumberAttribute,
     UnicodeAttribute, UnicodeSetAttribute, UTCDateTimeAttribute, BooleanAttribute, MapAttribute,
     ListAttribute, JSONAttribute, TTLAttribute, VersionAttribute)
 from pynamodb.constants import (
@@ -882,6 +882,37 @@ class OuterMapAttribute(MapAttribute):
         assert serialized == {'inner_map': {'M': {'foo': {'S': 'bar'}}}}
 
 
+class TestDynamicMapAttribute:
+
+    class CreatedAtTestModel(Model):
+        class CreatedAtMap(DynamicMapAttribute):
+            created_at = UTCDateTimeAttribute()
+        test_map = CreatedAtMap(default=dict)
+
+    def test_serialize(self):
+        test_model = TestDynamicMapAttribute.CreatedAtTestModel()
+        test_model.test_map.created_at = datetime(2017, 1, 1, tzinfo=timezone.utc)
+        test_model.test_map.foo = 'bar'
+        test_model.test_map.empty = None
+        assert test_model.serialize() == {'test_map': {'M': {
+            'created_at': {'S': '2017-01-01T00:00:00.000000+0000'},
+            'foo': {'S': 'bar'},
+            'empty': {'NULL': True},
+        }}}
+
+    def test_deserialize(self):
+        serialized = {'test_map': {'M': {
+            'created_at': {'S': '2017-01-01T00:00:00.000000+0000'},
+            'foo': {'S': 'bar'},
+            'empty': {'NULL': True},
+        }}}
+        test_model = TestDynamicMapAttribute.CreatedAtTestModel()
+        test_model.deserialize(serialized)
+        assert test_model.test_map.created_at == datetime(2017, 1, 1, tzinfo=timezone.utc)
+        assert test_model.test_map.foo == 'bar'
+        assert test_model.test_map.empty is None
+
+
 class TestListAttribute:
 
     def test_untyped_list(self):

tests/test_discriminator.py

@@ -1,6 +1,7 @@
 import pytest
 
 from pynamodb.attributes import DiscriminatorAttribute
+from pynamodb.attributes import DynamicMapAttribute
 from pynamodb.attributes import ListAttribute
 from pynamodb.attributes import MapAttribute
 from pynamodb.attributes import NumberAttribute
@@ -42,6 +43,23 @@ class ChildModel(DiscriminatorTestModel, discriminator='Child'):
     value = UnicodeAttribute()
 
 
+class DynamicSubclassedMapAttribute(DynamicMapAttribute):
+    string_attr = UnicodeAttribute()
+
+
+class DynamicMapDiscriminatorTestModel(Model, discriminator='Parent'):
+    class Meta:
+        host = 'http://localhost:8000'
+        table_name = 'test'
+    hash_key = UnicodeAttribute(hash_key=True)
+    value = DynamicSubclassedMapAttribute(default=dict)
+    type = DiscriminatorAttribute()
+
+
+class DynamicMapDiscriminatorChildTestModel(DynamicMapDiscriminatorTestModel, discriminator='Child'):
+    value = UnicodeAttribute()
+
+
 class TestDiscriminatorAttribute:
 
     def test_serialize(self):
@@ -128,3 +146,51 @@ def test_deserialize(self):
         assert isinstance(cm, ChildModel)
         assert cm.hash_key == 'foo'
         assert cm.value == 'bar'
+
+
+class TestDynamicDiscriminatorModel:
+
+    def test_serialize_parent(self):
+        m = DynamicMapDiscriminatorTestModel()
+        m.hash_key = 'foo'
+        m.value.string_attr = 'foostr'
+        m.value.bar_attribute = 3
+        assert m.serialize() == {
+            'hash_key': {'S': 'foo'},
+            'type': {'S': 'Parent'},
+            'value': {'M': {'string_attr': {'S': 'foostr'}, 'bar_attribute': {'N': '3'}}},
+        }
+
+    def test_deserialize_parent(self):
+        item = {
+            'hash_key': {'S': 'foo'},
+            'type': {'S': 'Parent'},
+            'value': {
+                'M': {'string_attr': {'S': 'foostr'}, 'bar_attribute': {'N': '3'}}
+            }
+        }
+        m = DynamicMapDiscriminatorTestModel.from_raw_data(item)
+        assert m.hash_key == 'foo'
+        assert m.value
+        assert m.value.string_attr == 'foostr'
+        assert m.value.bar_attribute == 3
+
+    def test_serialize_child(self):
+        m = DynamicMapDiscriminatorChildTestModel()
+        m.hash_key = 'foo'
+        m.value = 'string val'
+        assert m.serialize() == {
+            'hash_key': {'S': 'foo'},
+            'type': {'S': 'Child'},
+            'value': {'S': 'string val'}
+        }
+
+    def test_deserialize_child(self):
+        item = {
+            'hash_key': {'S': 'foo'},
+            'type': {'S': 'Child'},
+            'value': {'S': 'string val'}
+        }
+        m = DynamicMapDiscriminatorChildTestModel.from_raw_data(item)
+        assert m.hash_key == 'foo'
+        assert m.value == 'string val'