Support MapAttribute polymorphism using discriminators. (#836)

Author: jpinner-lyft

docs/index.rst

@@ -11,7 +11,6 @@ Features
 ========
 
 * Python 3 support
-* Python 2 support
 * Support for Unicode, Binary, JSON, Number, Set, and UTC Datetime attributes
 * Support for DynamoDB Local
 * Support for all of the DynamoDB API
@@ -32,6 +31,7 @@ Topics
    batch
    updates
    conditional
+   polymorphism
    attributes
    transaction
    optimistic_locking

docs/polymorphism.rst

@@ -0,0 +1,49 @@
+Polymorphism
+============
+
+PynamoDB supports polymorphism through the use of discriminators.
+
+A discriminator is a value that is written to DynamoDB that identifies the python class being stored.
+(Note: currently discriminators are only supported on MapAttribute subclasses; support for model subclasses coming soon.)
+
+Discriminator Attributes
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+The discriminator value is stored using a special attribute, the DiscriminatorAttribute.
+Only a single DiscriminatorAttribute can be defined on a class.
+
+The discriminator value can be assigned to a class as part of the definition:
+
+.. code-block:: python
+
+    class ParentClass(MapAttribute):
+        cls = DiscriminatorAttribute()
+
+    class ChildClass(ParentClass, discriminator='child'):
+        pass
+
+Declaring the discriminator value as part of the class definition will automatically register the class with the discriminator attribute.
+A class can also be registered manually:
+
+.. code-block:: python
+
+    class ParentClass(MapAttribute):
+        cls = DiscriminatorAttribute()
+
+    class ChildClass(ParentClass):
+        pass
+
+    ParentClass._cls.register_class(ChildClass, 'child')
+
+.. note::
+
+    A class may be registered with a discriminator attribute multiple times.
+    Only the first registered value is used during serialization;
+    however, any registered value can be used to deserialize the class.
+    This behavior is intended to facilitate migrations if discriminator values must be changed.
+
+.. warning::
+
+    Discriminator values are written to DynamoDB.
+    Changing the value after items have been saved to the database can result in deserialization failures.
+    In order to read items with an old discriminator value, the old value must be manually registered.

docs/release_notes.rst

@@ -1,9 +1,21 @@
 Release Notes
 =============
 
-v5.0 (unreleased)
------------------
+v5.0.0 (unreleased)
+-------------------
 
+:date: 2020-xx-xx
+
+This is major release and contains breaking changes. Please read the notes below carefully.
+
+**Polymorphism**
+
+This release introduces polymorphism support via ``DiscriminatorAttribute``.
+Discriminator values are written to DynamoDB and used during deserialization to instantiate the desired class.
+
+Other changes in this release:
+
+* Python 2 is no longer supported. Python 3.6 or greater is now required.
 * ``Model.query`` no longer demotes invalid range key conditions to be filter conditions to avoid surprising behaviors:
   where what's intended to be a cheap and fast condition ends up being expensive and slow. Since filter conditions
   cannot contain range keys, this had limited utility to begin with, and would sometimes cause confusing

pynamodb/attributes.py

@@ -13,7 +13,7 @@
 from dateutil.tz import tzutc
 from inspect import getfullargspec
 from inspect import getmembers
-from typing import Any, Callable, Dict, Generic, List, Mapping, Optional, TypeVar, Type, Union, Set, overload
+from typing import Any, Callable, Dict, Generic, List, Mapping, Optional, TypeVar, Type, Union, Set, cast, overload
 from typing import TYPE_CHECKING
 
 from pynamodb._compat import GenericMeta
@@ -218,12 +218,16 @@ def delete(self, *values: Any) -> 'DeleteAction':
 
 class AttributeContainerMeta(GenericMeta):
 
-    def __init__(self, name, bases, attrs, *args, **kwargs):
-        super().__init__(name, bases, attrs, *args, **kwargs)  # type: ignore
-        AttributeContainerMeta._initialize_attributes(self)
+    def __new__(cls, name, bases, namespace, discriminator=None):
+        # Defined so that the discriminator can be set in the class definition.
+        return super().__new__(cls, name, bases, namespace)
+
+    def __init__(self, name, bases, namespace, discriminator=None):
+        super().__init__(name, bases, namespace)
+        AttributeContainerMeta._initialize_attributes(self, discriminator)
 
     @staticmethod
-    def _initialize_attributes(cls):
+    def _initialize_attributes(cls, discriminator_value):
         """
         Initialize attributes on the class.
         """
@@ -249,6 +253,20 @@ def _initialize_attributes(cls):
                 # Prepend the `attr_path` lists with the dynamo attribute name.
                 attribute._update_attribute_paths(attribute.attr_name)
 
+        # Register the class with the discriminator if necessary.
+        discriminators = [name for name, attr in cls._attributes.items() if isinstance(attr, DiscriminatorAttribute)]
+        if len(discriminators) > 1:
+            raise ValueError("{} has more than one discriminator attribute: {}".format(
+                cls.__name__, ", ".join(discriminators)))
+        cls._discriminator = discriminators[0] if discriminators else None
+        # TODO(jpinner) add support for model polymorphism
+        if cls._discriminator and not issubclass(cls, MapAttribute):
+            raise NotImplementedError("Discriminators are not yet supported in model classes.")
+        if discriminator_value is not None:
+            if not cls._discriminator:
+                raise ValueError("{} does not have a discriminator attribute".format(cls.__name__))
+            cls._attributes[cls._discriminator].register_class(cls, discriminator_value)
+
 
 class AttributeContainer(metaclass=AttributeContainerMeta):
 
@@ -259,6 +277,7 @@ def __init__(self, _user_instantiated: bool = True, **attributes: Attribute) ->
         # instances do not have any Attributes defined and instead use this dictionary to store their
         # collection of name-value pairs.
         self.attribute_values: Dict[str, Any] = {}
+        self._set_discriminator()
         self._set_defaults(_user_instantiated=_user_instantiated)
         self._set_attributes(**attributes)
 
@@ -288,6 +307,15 @@ def _dynamo_to_python_attr(cls, dynamo_key: str) -> str:
         """
         return cls._dynamo_to_python_attrs.get(dynamo_key, dynamo_key)  # type: ignore
 
+    @classmethod
+    def _get_discriminator_attribute(cls) -> Optional['DiscriminatorAttribute']:
+        return cls.get_attributes()[cls._discriminator] if cls._discriminator else None  # type: ignore
+
+    def _set_discriminator(self) -> None:
+        discriminator_attr = self._get_discriminator_attribute()
+        if discriminator_attr and discriminator_attr.get_discriminator(self.__class__) is not None:
+            self.attribute_values[self._discriminator] = self.__class__  # type: ignore
+
     def _set_defaults(self, _user_instantiated: bool = True) -> None:
         """
         Sets and fields that provide a default value
@@ -336,6 +364,7 @@ def _deserialize(self, attribute_values: Dict[str, Dict[str, Any]]) -> None:
         Sets attributes sent back from DynamoDB on this object
         """
         self.attribute_values = {}
+        self._set_discriminator()
         self._set_defaults(_user_instantiated=False)
         for name, attr in self.get_attributes().items():
             attribute_value = attribute_values.get(attr.attr_name)
@@ -352,6 +381,47 @@ def __ne__(self, other: Any) -> bool:
         return self is not other
 
 
+class DiscriminatorAttribute(Attribute[type]):
+    attr_type = STRING
+
+    def __init__(self, attr_name: Optional[str] = None) -> None:
+        super().__init__(attr_name=attr_name)
+        self._class_map: Dict[type, Any] = {}
+        self._discriminator_map: Dict[Any, type] = {}
+
+    def register_class(self, cls: type, discriminator: Any):
+        discriminator = discriminator(cls) if callable(discriminator) else discriminator
+        current_class = self._discriminator_map.get(discriminator)
+        if current_class and current_class != cls:
+            raise ValueError("The discriminator value '{}' is already assigned to a class: {}".format(
+                discriminator, current_class.__name__))
+
+        if cls not in self._class_map:
+            self._class_map[cls] = discriminator
+
+        self._discriminator_map[discriminator] = cls
+
+    def get_discriminator(self, cls: type) -> Optional[Any]:
+        return self._class_map.get(cls)
+
+    def __set__(self, instance: Any, value: Optional[type]) -> None:
+        raise TypeError("'{}' object does not support item assignment".format(self.__class__.__name__))
+
+    def serialize(self, value):
+        """
+        Returns the discriminator value corresponding to the given class.
+        """
+        return self._class_map[value]
+
+    def deserialize(self, value):
+        """
+        Returns the class corresponding to the given discriminator value.
+        """
+        if value not in self._discriminator_map:
+            raise ValueError("Unknown discriminator value: {}".format(value))
+        return self._discriminator_map[value]
+
+
 class BinaryAttribute(Attribute[bytes]):
     """
     A binary attribute
@@ -861,7 +931,14 @@ def deserialize(self, values):
         """
         if not self.is_raw():
             # If this is a subclass of a MapAttribute (i.e typed), instantiate an instance
-            instance = type(self)()
+            cls = type(self)
+            discriminator_attr = cls._get_discriminator_attribute()
+            if discriminator_attr:
+                discriminator_attribute_value = values.pop(discriminator_attr.attr_name, None)
+                if discriminator_attribute_value:
+                    discriminator_value = discriminator_attr.get_value(discriminator_attribute_value)
+                    cls = discriminator_attr.deserialize(discriminator_value)
+            instance = cls()
             instance._deserialize(values)
             return instance
 

tests/test_discriminator.py

@@ -0,0 +1,105 @@
+import pytest
+
+from pynamodb.attributes import DiscriminatorAttribute
+from pynamodb.attributes import ListAttribute
+from pynamodb.attributes import MapAttribute
+from pynamodb.attributes import NumberAttribute
+from pynamodb.attributes import UnicodeAttribute
+from pynamodb.models import Model
+
+
+class_name = lambda cls: cls.__name__
+
+
+class TypedValue(MapAttribute):
+    _cls = DiscriminatorAttribute(attr_name = 'cls')
+    name = UnicodeAttribute()
+
+
+class NumberValue(TypedValue, discriminator=class_name):
+    value = NumberAttribute()
+
+
+class StringValue(TypedValue, discriminator=class_name):
+    value = UnicodeAttribute()
+
+
+class RenamedValue(TypedValue, discriminator='custom_name'):
+    value = UnicodeAttribute()
+
+
+class DiscriminatorTestModel(Model):
+    class Meta:
+        host = 'http://localhost:8000'
+        table_name = 'test'
+    hash_key = UnicodeAttribute(hash_key=True)
+    value = TypedValue()
+    values = ListAttribute(of=TypedValue)
+
+
+class TestDiscriminatorAttribute:
+
+    def test_serialize(self):
+        dtm = DiscriminatorTestModel()
+        dtm.hash_key = 'foo'
+        dtm.value = StringValue(name='foo', value='Hello')
+        dtm.values = [NumberValue(name='bar', value=5), RenamedValue(name='baz', value='World')]
+        assert dtm._serialize() == {
+            'HASH': 'foo',
+            'attributes': {
+                'value': {'M': {'cls': {'S': 'StringValue'}, 'name': {'S': 'foo'}, 'value': {'S': 'Hello'}}},
+                'values': {'L': [
+                    {'M': {'cls': {'S': 'NumberValue'}, 'name': {'S': 'bar'}, 'value': {'N': '5'}}},
+                    {'M': {'cls': {'S': 'custom_name'}, 'name': {'S': 'baz'}, 'value': {'S': 'World'}}}
+                ]}
+            }
+        }
+
+    def test_deserialize(self):
+        item = {
+            'hash_key': {'S': 'foo'},
+            'value': {'M': {'cls': {'S': 'StringValue'}, 'name': {'S': 'foo'}, 'value': {'S': 'Hello'}}},
+            'values': {'L': [
+                {'M': {'cls': {'S': 'NumberValue'}, 'name': {'S': 'bar'}, 'value': {'N': '5'}}},
+                {'M': {'cls': {'S': 'custom_name'}, 'name': {'S': 'baz'}, 'value': {'S': 'World'}}}
+            ]}
+        }
+        dtm = DiscriminatorTestModel.from_raw_data(item)
+        assert dtm.hash_key == 'foo'
+        assert dtm.value.value == 'Hello'
+        assert dtm.values[0].value == 5
+        assert dtm.values[1].value == 'World'
+
+    def test_condition_expression(self):
+        condition = DiscriminatorTestModel.value._cls == RenamedValue
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = condition.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0.#1 = :0"
+        assert placeholder_names == {'value': '#0', 'cls': '#1'}
+        assert expression_attribute_values == {':0': {'S': 'custom_name'}}
+
+    def test_multiple_discriminator_values(self):
+        class TestAttribute(MapAttribute, discriminator='new_value'):
+            cls = DiscriminatorAttribute()
+
+        TestAttribute.cls.register_class(TestAttribute, 'old_value')
+
+        # ensure the first registered value is used during serialization
+        assert TestAttribute.cls.get_discriminator(TestAttribute) == 'new_value'
+        assert TestAttribute.cls.serialize(TestAttribute) == 'new_value'
+
+        # ensure the second registered value can be used to deserialize
+        assert TestAttribute.cls.deserialize('old_value') == TestAttribute
+        assert TestAttribute.cls.deserialize('new_value') == TestAttribute
+
+    def test_multiple_discriminator_classes(self):
+        with pytest.raises(ValueError):
+            # fail when attempting to register a class with an existing discriminator value
+            class RenamedValue2(TypedValue, discriminator='custom_name'):
+                pass
+
+    def test_model(self):
+        with pytest.raises(NotImplementedError):
+            class DiscriminatedModel(Model):
+                hash_key = UnicodeAttribute(hash_key=True)
+                _cls = DiscriminatorAttribute()