Support document paths and attribute names containing dot characters. (#322)

Author: jpinner-lyft

pynamodb/attributes.py

@@ -73,33 +73,49 @@ def get_value(self, value):
     def __eq__(self, other):
         if other is None or isinstance(other, Attribute):  # handle object identity comparison
             return self is other
-        return Path(self.attr_name).__eq__(self._get_attribute_value(other))
+        return AttributePath(self).__eq__(other)
 
     def __lt__(self, other):
-        return Path(self.attr_name).__lt__(self._get_attribute_value(other))
+        return AttributePath(self).__lt__(other)
 
     def __le__(self, other):
-        return Path(self.attr_name).__le__(self._get_attribute_value(other))
+        return AttributePath(self).__le__(other)
 
     def __gt__(self, other):
-        return Path(self.attr_name).__gt__(self._get_attribute_value(other))
+        return AttributePath(self).__gt__(other)
 
     def __ge__(self, other):
-        return Path(self.attr_name).__ge__(self._get_attribute_value(other))
+        return AttributePath(self).__ge__(other)
 
-    def __getitem__(self, idx):  # support accessing list elements in condition expressions
-        if not isinstance(idx, int):
-            raise TypeError('list indices must be integers, not {}'.format(type(idx).__name__))
-        return Path('{0}[{1}]'.format(self.attr_name, idx))  # TODO include attribute value formatting
+    def __getitem__(self, idx):
+        return AttributePath(self)[idx]
 
-    def between(self, value1, value2):
-        return Path(self.attr_name).between(self._get_attribute_value(value1), self._get_attribute_value(value2))
+    def between(self, lower, upper):
+        return AttributePath(self).between(lower, upper)
 
     def startswith(self, prefix):
-        return Path(self.attr_name).startswith(self._get_attribute_value(prefix))
+        return AttributePath(self).startswith(prefix)
 
-    def _get_attribute_value(self, value):
-        return {ATTR_TYPE_MAP[self.attr_type]: self.serialize(value)}
+
+class AttributePath(Path):
+
+    def __init__(self, attribute):
+        super(AttributePath, self).__init__(attribute.attr_name, attribute_name=True)
+        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__))
+        return super(AttributePath, self).__getitem__(idx)
+
+    def _serialize(self, value):
+        if self.attribute.attr_type == LIST and not isinstance(value, list):
+            # List attributes assume the values to be serialized are lists.
+            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 {ATTR_TYPE_MAP[self.attribute.attr_type]: self.attribute.serialize(value)}
 
 
 class AttributeContainer(object):

pynamodb/expressions/condition.py

@@ -1,10 +1,21 @@
+from copy import copy
+from pynamodb.constants import AND, BETWEEN
 from pynamodb.expressions.util import get_value_placeholder, substitute_names
 
 
 class Path(object):
 
-    def __init__(self, path):
+    def __init__(self, path, attribute_name=False):
         self.path = path
+        self.attribute_name = attribute_name
+
+    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)
+        return element_path
 
     def __eq__(self, other):
         return self._compare('=', other)
@@ -21,18 +32,32 @@ def __gt__(self, other):
     def __ge__(self, other):
         return self._compare('>=', other)
 
-    def _compare(self, operator, value):
-        return Condition(self.path, operator, value)
+    def _compare(self, operator, other):
+        return Condition(self, operator, self._serialize(other))
 
-    def between(self, value1, value2):
+    def between(self, lower, upper):
         # This seemed preferable to other options such as merging value1 <= attribute & attribute <= value2
         # into one condition expression. DynamoDB only allows a single sort key comparison and having this
         # work but similar expressions like value1 <= attribute & attribute < value2 fail seems too brittle.
-        return Between(self.path, value1, value2)
+        return Between(self, self._serialize(lower), self._serialize(upper))
 
     def startswith(self, prefix):
         # A 'pythonic' replacement for begins_with to match string behavior (e.g. "foo".startswith("f"))
-        return BeginsWith(self.path, prefix)
+        return BeginsWith(self, self._serialize(prefix))
+
+    def _serialize(self, value):
+        # Allow subclasses to define value serialization.
+        return value
+
+    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
+
+    def __repr__(self):
+        return "Path('{0}', attribute_name={1})".format(self.path, self.attribute_name)
 
 
 class Condition(object):
@@ -46,7 +71,8 @@ def __init__(self, path, operator, *values):
         self.other_condition = None
 
     def serialize(self, placeholder_names, expression_attribute_values):
-        path = substitute_names(self.path, placeholder_names)
+        split = not self.path.attribute_name
+        path = substitute_names(self.path.path, placeholder_names, split=split)
         values = [get_value_placeholder(value, expression_attribute_values) for value in self.values]
         condition = self.format_string.format(*values, path=path, operator=self.operator)
         if self.logical_operator:
@@ -58,10 +84,18 @@ def __and__(self, other):
         if not isinstance(other, Condition):
             raise TypeError("unsupported operand type(s) for &: '{0}' and '{1}'",
                             self.__class__.__name__, other.__class__.__name__)
-        self.logical_operator = 'AND'
+        self.logical_operator = AND
         self.other_condition = other
         return self
 
+    def __repr__(self):
+        values = [value.items()[0][1] for value in self.values]
+        condition = self.format_string.format(*values, path=self.path, operator = self.operator)
+        if self.logical_operator:
+            other_conditions = repr(self.other_condition)
+            return '{0} {1} {2}'.format(condition, self.logical_operator, other_conditions)
+        return condition
+
     def __nonzero__(self):
         # Prevent users from accidentally comparing the condition object instead of the attribute instance
         raise TypeError("unsupported operand type(s) for bool: '{0}'".format(self.__class__.__name__))
@@ -74,12 +108,12 @@ def __bool__(self):
 class Between(Condition):
     format_string = '{path} {operator} {0} AND {1}'
 
-    def __init__(self, attribute, value1, value2):
-        super(Between, self).__init__(attribute, 'BETWEEN', value1, value2)
+    def __init__(self, path, lower, upper):
+        super(Between, self).__init__(path, BETWEEN, lower, upper)
 
 
 class BeginsWith(Condition):
     format_string = '{operator} ({path}, {0})'
 
-    def __init__(self, attribute, prefix):
-        super(BeginsWith, self).__init__(attribute, 'begins_with', prefix)
+    def __init__(self, path, prefix):
+        super(BeginsWith, self).__init__(path, 'begins_with', prefix)

pynamodb/expressions/projection.py

@@ -1,6 +1,19 @@
+from pynamodb.attributes import Attribute
+from pynamodb.expressions.condition import Path
 from pynamodb.expressions.util import substitute_names
 
 
 def create_projection_expression(attributes_to_get, placeholders):
-    expressions = [substitute_names(attribute, placeholders) for attribute in attributes_to_get]
+    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]
     return ', '.join(expressions)
+
+
+def _get_expression_split_pair(attribute):
+    if isinstance(attribute, Attribute):
+        return attribute.attr_name, False
+    if isinstance(attribute, Path):
+        return attribute.path, not attribute.attribute_name
+    return attribute, True

pynamodb/expressions/util.py

@@ -3,12 +3,12 @@
 PATH_SEGMENT_REGEX = re.compile(r'([^\[\]]+)((?:\[\d+\])*)$')
 
 
-def substitute_names(expression, placeholders):
+def substitute_names(expression, placeholders, split=True):
     """
     Replaces names in the given expression with placeholders.
     Stores the placeholders in the given dictionary.
     """
-    path_segments = expression.split('.')
+    path_segments = expression.split('.') if split else [expression]
     for idx, segment in enumerate(path_segments):
         match = PATH_SEGMENT_REGEX.match(segment)
         if not match:

pynamodb/tests/test_expressions.py

@@ -1,8 +1,36 @@
-from pynamodb.attributes import UnicodeAttribute
+from pynamodb.attributes import ListAttribute, UnicodeAttribute
 from pynamodb.compat import CompatTestCase as TestCase
+from pynamodb.expressions.condition import Path
 from pynamodb.expressions.projection import create_projection_expression
 
 
+class PathTestCase(TestCase):
+
+    def test_document_path(self):
+        path = Path('foo.bar')
+        assert str(path) == 'foo.bar'
+        assert repr(path) == "Path('foo.bar', attribute_name=False)"
+
+    def test_attribute_name(self):
+        path = Path('foo.bar', attribute_name=True)
+        assert str(path) == "'foo.bar'"
+        assert repr(path) == "Path('foo.bar', attribute_name=True)"
+
+    def test_index_document_path(self):
+        path = Path('foo.bar')[0]
+        assert str(path) == 'foo.bar[0]'
+        assert repr(path) == "Path('foo.bar[0]', attribute_name=False)"
+
+    def test_index_attribute_name(self):
+        path = Path('foo.bar', attribute_name=True)[0]
+        assert str(path) == "'foo.bar'[0]"
+        assert repr(path) == "Path('foo.bar[0]', attribute_name=True)"
+
+    def test_index_invalid(self):
+        with self.assertRaises(TypeError):
+            Path('foo.bar')['foo']
+
+
 class ProjectionExpressionTestCase(TestCase):
 
     def test_create_projection_expression(self):
@@ -25,6 +53,42 @@ def test_create_projection_expression_invalid_attribute_raises(self):
             with self.assertRaises(ValueError):
                 create_projection_expression([attribute], {})
 
+    def test_create_project_expression_with_document_paths(self):
+        attributes_to_get = [Path('foo.bar')[0]]
+        placeholders = {}
+        projection_expression = create_projection_expression(attributes_to_get, placeholders)
+        assert projection_expression == "#0.#1[0]"
+        assert placeholders == {'foo': '#0', 'bar': '#1'}
+
+    def test_create_project_expression_with_attribute_names(self):
+        attributes_to_get = [Path('foo.bar', attribute_name=True)[0]]
+        placeholders = {}
+        projection_expression = create_projection_expression(attributes_to_get, placeholders)
+        assert projection_expression == "#0[0]"
+        assert placeholders == {'foo.bar': '#0'}
+
+    def test_create_projection_expression_with_attributes(self):
+        attributes_to_get = [
+            UnicodeAttribute(attr_name='ProductReviews.FiveStar'),
+            UnicodeAttribute(attr_name='ProductReviews.ThreeStar'),
+            UnicodeAttribute(attr_name='ProductReviews.OneStar')
+        ]
+        placeholders = {}
+        projection_expression = create_projection_expression(attributes_to_get, placeholders)
+        assert projection_expression == "#0, #1, #2"
+        assert placeholders == {
+            'ProductReviews.FiveStar': '#0',
+            'ProductReviews.ThreeStar': '#1',
+            'ProductReviews.OneStar': '#2',
+        }
+
+    def test_create_projection_expression_not_a_list(self):
+        attributes_to_get = 'Description'
+        placeholders = {}
+        projection_expression = create_projection_expression(attributes_to_get, placeholders)
+        assert projection_expression == "#0"
+        assert placeholders == {'Description': '#0'}
+
 
 class ConditionExpressionTestCase(TestCase):
 
@@ -88,9 +152,38 @@ def test_begins_with(self):
         assert expression_attribute_values == {':0': {'S' : 'bar'}}
 
     def test_indexing(self):
-        condition = self.attribute[0] == 'bar'
+        condition = ListAttribute(attr_name='foo')[0] == 'bar'
         placeholder_names, expression_attribute_values = {}, {}
         expression = condition.serialize(placeholder_names, expression_attribute_values)
         assert expression == "#0[0] = :0"
         assert placeholder_names == {'foo': '#0'}
-        assert expression_attribute_values == {':0': 'bar'}  # TODO fix attribute value formatting
+        assert expression_attribute_values == {':0': {'S' : 'bar'}}
+
+    def test_invalid_indexing(self):
+        with self.assertRaises(TypeError):
+            self.attribute[0]
+
+    def test_double_indexing(self):
+        condition = ListAttribute(attr_name='foo')[0][1] == 'bar'
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = condition.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0[0][1] = :0"
+        assert placeholder_names == {'foo': '#0'}
+        assert expression_attribute_values == {':0': {'S' : 'bar'}}
+
+    def test_list_comparison(self):
+        condition = ListAttribute(attr_name='foo') == ['bar', 'baz']
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = condition.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0 = :0"
+        assert placeholder_names == {'foo': '#0'}
+        assert expression_attribute_values == {':0': {'L': [{'S' : 'bar'}, {'S': 'baz'}]}}
+
+    def test_dotted_attribute_name(self):
+        self.attribute.attr_name = 'foo.bar'
+        condition = self.attribute == 'baz'
+        placeholder_names, expression_attribute_values = {}, {}
+        expression = condition.serialize(placeholder_names, expression_attribute_values)
+        assert expression == "#0 = :0"
+        assert placeholder_names == {'foo.bar': '#0'}
+        assert expression_attribute_values == {':0': {'S': 'baz'}}