Support MapAttribute map dereferencing in condition expressions. (#343)

Author: jpinner-lyft

pynamodb/attributes.py

@@ -5,6 +5,7 @@
 from six import add_metaclass
 import json
 from base64 import b64encode, b64decode
+from copy import deepcopy
 from datetime import datetime
 from dateutil.parser import parse
 from dateutil.tz import tzutc
@@ -36,15 +37,27 @@ def __init__(self,
             self.null = null
         self.is_hash_key = hash_key
         self.is_range_key = range_key
-        self.attr_name = attr_name
+        self.attr_path = [attr_name]
+
+    @property
+    def attr_name(self):
+        return self.attr_path[-1]
+
+    @attr_name.setter
+    def attr_name(self, value):
+        self.attr_path[-1] = value
 
     def __set__(self, instance, value):
         if instance and not self._is_map_attribute_class_object(instance):
             attr_name = instance._dynamo_to_python_attrs.get(self.attr_name, self.attr_name)
             instance.attribute_values[attr_name] = value
 
     def __get__(self, instance, owner):
-        if instance and not self._is_map_attribute_class_object(instance):
+        if self._is_map_attribute_class_object(instance):
+            # MapAttribute class objects store a local copy of the attribute with `attr_path` set to the document path.
+            attr_name = instance._dynamo_to_python_attrs.get(self.attr_name, self.attr_name)
+            return instance.__dict__.get(attr_name, None) or self
+        elif instance:
             attr_name = instance._dynamo_to_python_attrs.get(self.attr_name, self.attr_name)
             return instance.attribute_values.get(attr_name, None)
         else:
@@ -121,12 +134,14 @@ def contains(self, item):
 class AttributePath(Path):
 
     def __init__(self, attribute):
-        super(AttributePath, self).__init__(attribute.attr_name, attribute_name=True)
+        super(AttributePath, self).__init__(attribute.attr_path)
         self.attribute = attribute
 
     def __getitem__(self, idx):
         if self.attribute.attr_type != LIST:  # only list elements support the list dereference operator
             raise TypeError("'{0}' object has no attribute __getitem__".format(self.attribute.__class__.__name__))
+        # The __getitem__ call returns a new Path instance, not an AttributePath instance.
+        # This is intended since the list element is not the same attribute as the list itself.
         return super(AttributePath, self).__getitem__(idx)
 
     def contains(self, item):
@@ -145,7 +160,7 @@ def _serialize(self, value):
             return self.attribute.serialize([value])[0]
         if self.attribute.attr_type == MAP and not isinstance(value, dict):
             # Map attributes assume the values to be serialized are maps.
-            return self.attribute.serialize({'': value})['']
+            return super(AttributePath, self)._serialize(value)
         return {ATTR_TYPE_MAP[self.attribute.attr_type]: self.attribute.serialize(value)}
 
 
@@ -174,17 +189,24 @@ def _initialize_attributes(cls):
 
             if issubclass(item_cls, Attribute):
                 instance = getattr(cls, item_name)
+                initialized = False
                 if isinstance(instance, MapAttribute):
                     # MapAttribute instances that are class attributes of an AttributeContainer class
                     # should behave like an Attribute instance and not an AttributeContainer instance.
-                    instance._make_attribute()
+                    initialized = instance._make_attribute()
 
                 cls._attributes[item_name] = instance
                 if instance.attr_name is not None:
                     cls._dynamo_to_python_attrs[instance.attr_name] = item_name
                 else:
                     instance.attr_name = item_name
 
+                if initialized and isinstance(instance, MapAttribute):
+                    # To support creating expressions from nested attributes, MapAttribute instances
+                    # store local copies of the attributes in cls._attributes with `attr_path` set.
+                    # Prepend the `attr_path` lists with the dynamo attribute name.
+                    instance._update_attribute_paths(instance.attr_name)
+
 
 @add_metaclass(AttributeContainerMeta)
 class AttributeContainer(object):
@@ -599,7 +621,7 @@ def __init__(self, **attributes):
         # It is possible that attributes names can collide with argument names of Attribute.__init__.
         # Assume that this is the case if any of the following are true:
         #   - the user passed in other attributes that did not match any argument names
-        #   - this is "raw" (i.e. non-subclassed) MapAttribute instance and attempting to store the attributes
+        #   - this is a "raw" (i.e. non-subclassed) MapAttribute instance and attempting to store the attributes
         #     cannot raise a ValueError (if this assumption is wrong, calling `_make_attribute` removes them)
         #   - the names of all attributes in self.attribute_kwargs match attributes defined on the class
         if self.attribute_kwargs and (
@@ -615,12 +637,30 @@ def _is_attribute_container(self):
     def _make_attribute(self):
         # WARNING! This function is only intended to be called from the AttributeContainerMeta metaclass.
         if not self._is_attribute_container():
-            return
+            # This instance has already been initialized by another AttributeContainer class.
+            return False
         # During initialization the kwargs were stored in `attribute_kwargs`. Remove them and re-initialize the class.
         kwargs = self.attribute_kwargs
         del self.attribute_kwargs
         del self.attribute_values
         Attribute.__init__(self, **kwargs)
+        for name, attr in self._get_attributes().items():
+            # Set a local attribute with the same name that shadows the class attribute.
+            # Because attr is a data descriptor and the attribute already exists on the class,
+            # we have to store the local copy directly into __dict__ to prevent calling attr.__set__.
+            # Use deepcopy so that `attr_path` and any local attributes are also copied.
+            self.__dict__[name] = deepcopy(attr)
+        return True
+
+    def _update_attribute_paths(self, path_segment):
+        # WARNING! This function is only intended to be called from the AttributeContainerMeta metaclass.
+        if self._is_attribute_container():
+            raise AssertionError("MapAttribute._update_attribute_paths called before MapAttribute._make_attribute")
+        for name in self._get_attributes().keys():
+            local_attr = self.__dict__[name]
+            local_attr.attr_path.insert(0, path_segment)
+            if isinstance(local_attr, MapAttribute):
+                local_attr._update_attribute_paths(path_segment)
 
     def __iter__(self):
         return iter(self.attribute_values)

pynamodb/connection/base.py

@@ -1433,7 +1433,7 @@ def _get_condition(self, table_name, attribute_name, operator, *values):
             {self.get_attribute_type(table_name, attribute_name, value): self.parse_attribute(value)}
             for value in values
         ]
-        return getattr(Path(attribute_name, attribute_name=True), operator)(*values)
+        return getattr(Path([attribute_name]), operator)(*values)
 
     def _check_condition(self, name, condition, expected_or_filter, conditional_operator):
         if condition is not None:

pynamodb/expressions/condition.py

@@ -1,8 +1,8 @@
-from copy import copy
 from pynamodb.constants import (
     AND, ATTR_TYPE_MAP, BETWEEN, BINARY_SHORT, IN, NUMBER_SHORT, OR, SHORT_ATTR_TYPES, STRING_SHORT
 )
 from pynamodb.expressions.util import get_value_placeholder, substitute_names
+from six import string_types
 from six.moves import range
 
 
@@ -90,16 +90,17 @@ class Path(Operand):
     In addition to supporting comparisons, Path also supports creating conditions from functions.
     """
 
-    def __init__(self, path, attribute_name=False):
-        self.path = path
-        self.attribute_name = attribute_name
+    def __init__(self, path):
+        if not path:
+            raise ValueError("path cannot be empty")
+        self.path = path.split('.') if isinstance(path, string_types) else list(path)
 
     def __getitem__(self, idx):
         # list dereference operator
         if not isinstance(idx, int):
             raise TypeError("list indices must be integers, not {0}".format(type(idx).__name__))
-        element_path = copy(self)
-        element_path.path = '{0}[{1}]'.format(self.path, idx)
+        element_path = Path(self.path)  # copy the document path before indexing last element
+        element_path.path[-1] = '{0}[{1}]'.format(self.path[-1], idx)
         return element_path
 
     def exists(self):
@@ -119,14 +120,17 @@ def contains(self, item):
         return Contains(self, self._serialize(item))
 
     def __str__(self):
-        if self.attribute_name and '.' in self.path:
-            # Quote the path to illustrate that the dot characters are not dereference operators.
-            path, sep, rem = self.path.partition('[')
-            return repr(path) + sep + rem
-        return self.path
+        # Quote the path to illustrate that any dot characters are not dereference operators.
+        quoted_path = [self._quote_path(segment) if '.' in segment else segment for segment in self.path]
+        return '.'.join(quoted_path)
 
     def __repr__(self):
-        return "Path('{0}', attribute_name={1})".format(self.path, self.attribute_name)
+        return "Path({0})".format(self.path)
+
+    @staticmethod
+    def _quote_path(path):
+        path, sep, rem = path.partition('[')
+        return repr(path) + sep + rem
 
 
 class Condition(object):
@@ -148,8 +152,7 @@ def serialize(self, placeholder_names, expression_attribute_values):
 
     def _get_path(self, path, placeholder_names):
         if isinstance(path, Path):
-            split = not path.attribute_name
-            return substitute_names(path.path, placeholder_names, split=split)
+            return substitute_names(path.path, placeholder_names)
         elif isinstance(path, Size):
             return "size ({0})".format(self._get_path(path.path, placeholder_names))
         else:

pynamodb/expressions/projection.py

@@ -6,14 +6,13 @@
 def create_projection_expression(attributes_to_get, placeholders):
     if not isinstance(attributes_to_get, list):
         attributes_to_get = [attributes_to_get]
-    expression_split_pairs = [_get_expression_split_pair(attribute) for attribute in attributes_to_get]
-    expressions = [substitute_names(expr, placeholders, split=split) for (expr, split) in expression_split_pairs]
+    expressions = [substitute_names(_get_document_path(attribute), placeholders) for attribute in attributes_to_get]
     return ', '.join(expressions)
 
 
-def _get_expression_split_pair(attribute):
+def _get_document_path(attribute):
     if isinstance(attribute, Attribute):
-        return attribute.attr_name, False
+        return [attribute.attr_name]
     if isinstance(attribute, Path):
-        return attribute.path, not attribute.attribute_name
-    return attribute, True
+        return attribute.path
+    return attribute.split('.')

pynamodb/expressions/update.py

@@ -10,7 +10,7 @@ def __init__(self, path, value=None):
         self.value = value
 
     def serialize(self, placeholder_names, expression_attribute_values):
-        path = substitute_names(self.path, placeholder_names, split=True)
+        path = substitute_names(self.path.split('.'), placeholder_names)
         value = get_value_placeholder(self.value, expression_attribute_values) if self.value else None
         return self.format_string.format(value, path=path)
 

pynamodb/expressions/util.py

@@ -3,24 +3,30 @@
 PATH_SEGMENT_REGEX = re.compile(r'([^\[\]]+)((?:\[\d+\])*)$')
 
 
-def substitute_names(expression, placeholders, split=True):
+def substitute_names(document_path, placeholders):
     """
-    Replaces names in the given expression with placeholders.
+    Replaces all attribute names in the given document path with placeholders.
     Stores the placeholders in the given dictionary.
+
+    :param document_path: list of path segments (an attribute name and optional list dereference)
+    :param placeholders:  a dictionary to store mappings from attribute names to expression attribute name placeholders
+
+    For example: given the document_path for some attribute "baz", that is the first element of a list attribute "bar",
+    that itself is a map element of "foo" (i.e. ['foo', 'bar[0], 'baz']) and an empty placeholders dictionary,
+    `substitute_names` will return "#0.#1[0].#2" and placeholders will contain {"foo": "#0", "bar": "#1", "baz": "#2}
     """
-    path_segments = expression.split('.') if split else [expression]
-    for idx, segment in enumerate(path_segments):
+    for idx, segment in enumerate(document_path):
         match = PATH_SEGMENT_REGEX.match(segment)
         if not match:
-            raise ValueError('{0} is not a valid document path'.format(expression))
+            raise ValueError('{0} is not a valid document path'.format('.'.join(document_path)))
         name, indexes = match.groups()
         if name in placeholders:
             placeholder = placeholders[name]
         else:
             placeholder = '#' + str(len(placeholders))
             placeholders[name] = placeholder
-        path_segments[idx] = placeholder + indexes
-    return '.'.join(path_segments)
+        document_path[idx] = placeholder + indexes
+    return '.'.join(document_path)
 
 
 def get_value_placeholder(value, expression_attribute_values):

pynamodb/models.py

@@ -215,22 +215,22 @@ class Model(AttributeContainer):
     _throttle = NoThrottle()
     DoesNotExist = DoesNotExist
 
-    def __init__(self, hash_key=None, range_key=None, **attrs):
+    def __init__(self, hash_key=None, range_key=None, **attributes):
         """
         :param hash_key: Required. The hash key for this object.
         :param range_key: Only required if the table has a range key attribute.
         :param attrs: A dictionary of attributes to set on this object.
         """
         if hash_key is not None:
-            attrs[self._dynamo_to_python_attr(self._get_meta_data().hash_keyname)] = hash_key
+            attributes[self._dynamo_to_python_attr(self._get_meta_data().hash_keyname)] = hash_key
         if range_key is not None:
             range_keyname = self._get_meta_data().range_keyname
             if range_keyname is None:
                 raise ValueError(
                     "This table has no range key, but a range key value was provided: {0}".format(range_key)
                 )
-            attrs[self._dynamo_to_python_attr(range_keyname)] = range_key
-        super(Model, self).__init__(**attrs)
+            attributes[self._dynamo_to_python_attr(range_keyname)] = range_key
+        super(Model, self).__init__(**attributes)
 
     @classmethod
     def has_map_or_list_attributes(cls):

pynamodb/tests/test_attributes.py

@@ -730,6 +730,53 @@ class ThingModel(Model):
     def test_metaclass(self):
         assert type(MapAttribute) == MapAttributeMeta
 
+    def test_attribute_paths_subclassing(self):
+        class SubMapAttribute(MapAttribute):
+            foo = UnicodeAttribute(attr_name='dyn_foo')
+
+        class SubSubMapAttribute(SubMapAttribute):
+            bar = UnicodeAttribute(attr_name='dyn_bar')
+
+        class SubModel(Model):
+            sub_map = SubMapAttribute(attr_name='dyn_sub_map')
+
+        class SubSubModel(SubModel):
+            sub_sub_map = SubSubMapAttribute()
+
+        assert SubModel.sub_map.foo.attr_name == 'dyn_foo'
+        assert SubModel.sub_map.foo.attr_path == ['dyn_sub_map', 'dyn_foo']
+        assert SubSubModel.sub_map.foo.attr_name == 'dyn_foo'
+        assert SubSubModel.sub_map.foo.attr_path == ['dyn_sub_map', 'dyn_foo']
+        assert SubSubModel.sub_sub_map.foo.attr_name == 'dyn_foo'
+        assert SubSubModel.sub_sub_map.foo.attr_path == ['sub_sub_map', 'dyn_foo']
+        assert SubSubModel.sub_sub_map.bar.attr_name == 'dyn_bar'
+        assert SubSubModel.sub_sub_map.bar.attr_path == ['sub_sub_map', 'dyn_bar']
+
+    def test_attribute_paths_wrapping(self):
+        class InnerMapAttribute(MapAttribute):
+            map_attr = MapAttribute(attr_name='dyn_map_attr')
+
+        class MiddleMapAttributeA(MapAttribute):
+            inner_map = InnerMapAttribute(attr_name='dyn_in_map_a')
+
+        class MiddleMapAttributeB(MapAttribute):
+            inner_map = InnerMapAttribute(attr_name='dyn_in_map_b')
+
+        class OuterMapAttribute(MapAttribute):
+            mid_map_a = MiddleMapAttributeA()
+            mid_map_b = MiddleMapAttributeB()
+
+        class MyModel(Model):
+            outer_map = OuterMapAttribute(attr_name='dyn_out_map')
+
+        mid_map_a_map_attr = MyModel.outer_map.mid_map_a.inner_map.map_attr
+        mid_map_b_map_attr = MyModel.outer_map.mid_map_b.inner_map.map_attr
+
+        assert mid_map_a_map_attr.attr_name == 'dyn_map_attr'
+        assert mid_map_a_map_attr.attr_path == ['dyn_out_map', 'mid_map_a', 'dyn_in_map_a', 'dyn_map_attr']
+        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']
+
 
 class TestValueDeserialize:
     def test__get_value_for_deserialize(self):