Better handling of MapAttribute iteration and item assignment. (#367)

Author: jpinner-lyft

pynamodb/attributes.py

@@ -82,6 +82,10 @@ def get_value(self, value):
         serialized_dynamo_type = ATTR_TYPE_MAP[self.attr_type]
         return value.get(serialized_dynamo_type)
 
+    def __iter__(self):
+        # Because we define __getitem__ below for condition expression support
+        raise TypeError("'{0}' object is not iterable".format(self.__class__.__name__))
+
     # Condition Expression Support
     def __eq__(self, other):
         if other is None or isinstance(other, Attribute):  # handle object identity comparison
@@ -664,7 +668,9 @@ def _update_attribute_paths(self, path_segment):
                 local_attr._update_attribute_paths(path_segment)
 
     def __iter__(self):
-        return iter(self.attribute_values)
+        if self._is_attribute_container():
+            return iter(self.attribute_values)
+        return super(MapAttribute, self).__iter__()
 
     def __getitem__(self, item):
         if self._is_attribute_container():
@@ -679,6 +685,16 @@ def __getitem__(self, item):
         else:
             raise AttributeError("'{0}' has no attribute '{1}'".format(self.__class__.__name__, item))
 
+    def __setitem__(self, item, value):
+        if not self._is_attribute_container():
+            raise TypeError("'{0}' object does not support item assignment".format(self.__class__.__name__))
+        if self.is_raw():
+            self.attribute_values[item] = value
+        elif item in self._attributes:
+            setattr(self, item, value)
+        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.
         # MapAttribute subclasses should access attributes via the Attribute descriptors.

pynamodb/expressions/operand.py

@@ -241,6 +241,10 @@ def __init__(self, attribute_or_path):
     def path(self):
         return self.values[0]
 
+    def __iter__(self):
+        # Because we define __getitem__ Path is considered an iterable
+        raise TypeError("'{0}' object is not iterable".format(self.__class__.__name__))
+
     def __getitem__(self, item):
         # The __getitem__ call returns a new Path instance without any attribute set.
         # This is intended since the nested element is not the same attribute as ``self``.

pynamodb/tests/test_attributes.py

@@ -606,6 +606,17 @@ def test_raw_set_attr(self):
         assert item.map_attr['num'] == 3
         assert item.map_attr['nested']['nestedfoo'] == 'nestedbar'
 
+    def test_raw_set_item(self):
+        item = AttributeTestModel()
+        item.map_attr = {}
+        item.map_attr['foo'] = 'bar'
+        item.map_attr['num'] = 3
+        item.map_attr['nested'] = {'nestedfoo': 'nestedbar'}
+
+        assert item.map_attr['foo'] == 'bar'
+        assert item.map_attr['num'] == 3
+        assert item.map_attr['nested']['nestedfoo'] == 'nestedbar'
+
     def test_raw_map_from_dict(self):
         item = AttributeTestModel(
             map_attr={
@@ -633,6 +644,18 @@ def test_raw_map_access(self):
         for k, v in six.iteritems(raw):
             assert attr[k] == v
 
+    def test_raw_map_iter(self):
+        raw = {
+            "foo": "bar",
+            "num": 3,
+            "nested": {
+                "nestedfoo": "nestedbar"
+            }
+        }
+        attr = MapAttribute(**raw)
+
+        assert list(iter(raw)) == list(iter(attr))
+
     def test_raw_map_json_serialize(self):
         raw = {
             "foo": "bar",