Support "raw" MapAttribute dereferencing in condition expressions. (#362)

Author: jpinner-lyft

docs/conditional.rst

@@ -54,13 +54,24 @@ for more details.
     OR, \|, (Thread.views < 1) | (Thread.views > 5)
     NOT, ~, ~Thread.subject.contains('foobar')
 
-If necessary, you can use document paths to access nested list and map attributes:
+Conditions expressions using nested list and map attributes can be created with Python's item operator ``[]``:
 
 .. code-block:: python
 
-    from pynamodb.expressions.condition import size
+    from pynamodb.models import Model
+    from pynamodb.attributes import (
+        ListAttribute, MapAttribute, UnicodeAttribute
+    )
+
+    class Container(Model):
+        class Meta:
+            table_name = 'Container'
+
+        name = UnicodeAttribute(hash_key = True)
+        my_map = MapAttribute()
+        my_list = ListAttribute()
 
-    print(size('foo.bar[0].baz') == 0)
+    print(Container.my_map['foo'].exists() | Container.my_list[0].contains('bar'))
 
 
 Conditional Model.save

pynamodb/attributes.py

@@ -667,7 +667,17 @@ def __iter__(self):
         return iter(self.attribute_values)
 
     def __getitem__(self, item):
-        return self.attribute_values[item]
+        if self._is_attribute_container():
+            return self.attribute_values[item]
+        # 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)])
+        elif item in self._attributes:
+            return getattr(self, item)
+        else:
+            raise AttributeError("'{0}' has no attribute '{1}'".format(self.__class__.__name__, item))
 
     def __getattr__(self, attr):
         # This should only be called for "raw" (i.e. non-subclassed) MapAttribute instances.

pynamodb/expressions/operand.py

@@ -1,5 +1,6 @@
 from pynamodb.constants import (
-    ATTR_TYPE_MAP, BINARY_SET, LIST, LIST_SHORT, MAP, NUMBER_SET, NUMBER_SHORT, SHORT_ATTR_TYPES, STRING_SET, STRING_SHORT
+    ATTR_TYPE_MAP, BINARY_SET, LIST, LIST_SHORT, MAP, MAP_SHORT,
+    NUMBER_SET, NUMBER_SHORT, SHORT_ATTR_TYPES, STRING_SET, STRING_SHORT
 )
 from pynamodb.expressions.condition import (
     BeginsWith, Between, Comparison, Contains, Exists, In, IsType, NotExists
@@ -8,6 +9,7 @@
     AddAction, DeleteAction, RemoveAction, SetAction
 )
 from pynamodb.expressions.util import get_path_segments, get_value_placeholder, substitute_names
+from six import string_types
 
 
 class _Operand(object):
@@ -239,17 +241,24 @@ def __init__(self, attribute_or_path):
     def path(self):
         return self.values[0]
 
-    def __getitem__(self, idx):
-        # list dereference operator
-        if self.attribute and self.attribute.attr_type != LIST:
-            raise TypeError("'{0}' object has no attribute __getitem__".format(self.attribute.__class__.__name__))
-        if not isinstance(idx, int):
-            raise TypeError("list indices must be integers, not {0}".format(type(idx).__name__))
+    def __getitem__(self, item):
         # The __getitem__ call returns a new Path instance without any attribute set.
-        # This is intended since the list element is not the same attribute as the list itself.
-        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
+        # This is intended since the nested element is not the same attribute as ``self``.
+        if self.attribute and self.attribute.attr_type not in [LIST, MAP]:
+            raise TypeError("'{0}' object has no attribute __getitem__".format(self.attribute.__class__.__name__))
+        if self.short_attr_type == LIST_SHORT and not isinstance(item, int):
+            raise TypeError("list indices must be integers, not {0}".format(type(item).__name__))
+        if self.short_attr_type == MAP_SHORT and not isinstance(item, string_types):
+            raise TypeError("map attributes must be strings, not {0}".format(type(item).__name__))
+        if isinstance(item, int):
+            # list dereference operator
+            element_path = Path(self.path)  # copy the document path before indexing last element
+            element_path.path[-1] = '{0}[{1}]'.format(self.path[-1], item)
+            return element_path
+        if isinstance(item, string_types):
+            # map dereference operator
+            return Path(self.path + [item])
+        raise TypeError("item must be an integer or string, not {0}".format(type(item).__name__))
 
     def __or__(self, other):
         return _IfNotExists(self, self._to_operand(other))

pynamodb/tests/test_expressions.py

@@ -28,9 +28,14 @@ def test_index_attribute_name(self):
         assert str(path) == "'foo.bar'[0]"
         assert repr(path) == "Path(['foo.bar[0]'])"
 
+    def test_index_map_attribute(self):
+        path = Path(['foo.bar'])['baz']
+        assert str(path) == "'foo.bar'.baz"
+        assert repr(path) == "Path(['foo.bar', 'baz'])"
+
     def test_index_invalid(self):
         with self.assertRaises(TypeError):
-            Path('foo.bar')['foo']
+            Path('foo.bar')[0.0]
 
 
 class ActionTestCase(TestCase):
@@ -326,6 +331,19 @@ def test_dotted_attribute_name(self):
         assert placeholder_names == {'foo.bar': '#0'}
         assert expression_attribute_values == {':0': {'S': 'baz'}}
 
+    def test_map_attribute_indexing(self):
+        # Simulate initialization from inside an AttributeContainer
+        my_map_attribute = MapAttribute(attr_name='foo.bar')
+        my_map_attribute._make_attribute()
+        my_map_attribute._update_attribute_paths(my_map_attribute.attr_name)
+
+        condition = my_map_attribute['foo'] == 'baz'
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = condition.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0.#1 = :0"
+        assert placeholder_names == {'foo.bar': '#0', 'foo': '#1'}
+        assert expression_attribute_values == {':0': {'S': 'baz'}}
+
     def test_map_attribute_dereference(self):
         class MyMapAttribute(MapAttribute):
             nested_string = self.attribute
@@ -342,6 +360,34 @@ class MyMapAttribute(MapAttribute):
         assert placeholder_names == {'foo.bar': '#0', 'foo': '#1'}
         assert expression_attribute_values == {':0': {'S': 'baz'}}
 
+    def test_map_attribute_dereference_via_indexing(self):
+        class MyMapAttribute(MapAttribute):
+            nested_string = self.attribute
+
+        # Simulate initialization from inside an AttributeContainer
+        my_map_attribute = MyMapAttribute(attr_name='foo.bar')
+        my_map_attribute._make_attribute()
+        my_map_attribute._update_attribute_paths(my_map_attribute.attr_name)
+
+        condition = my_map_attribute['nested_string'] == 'baz'
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = condition.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0.#1 = :0"
+        assert placeholder_names == {'foo.bar': '#0', 'foo': '#1'}
+        assert expression_attribute_values == {':0': {'S': 'baz'}}
+
+    def test_map_attribute_dereference_via_indexing_missing_attribute(self):
+        class MyMapAttribute(MapAttribute):
+            nested_string = self.attribute
+
+        # Simulate initialization from inside an AttributeContainer
+        my_map_attribute = MyMapAttribute(attr_name='foo.bar')
+        my_map_attribute._make_attribute()
+        my_map_attribute._update_attribute_paths(my_map_attribute.attr_name)
+
+        with self.assertRaises(AttributeError):
+            my_map_attribute['missing_attribute'] == 'baz'
+
 
 class UpdateExpressionTestCase(TestCase):