Improve typing in expressions module. (#851)

Author: jpinner-lyft

pynamodb/connection/base.py

@@ -1145,7 +1145,7 @@ def batch_get_item(
             }
         }
 
-        args_map = {}
+        args_map: Dict[str, Any] = {}
         name_placeholders: Dict[str, str] = {}
         if consistent_read:
             args_map[CONSISTENT_READ] = consistent_read

pynamodb/constants.py

@@ -201,8 +201,3 @@
 META_CLASS_NAME = "Meta"
 REGION = "region"
 HOST = "host"
-
-AND = 'AND'
-OR = 'OR'
-BETWEEN = 'BETWEEN'
-IN = 'IN'

pynamodb/expressions/condition.py

@@ -1,4 +1,4 @@
-from pynamodb.constants import AND, BETWEEN, IN, OR
+from typing import Dict
 
 
 # match dynamo function syntax: size(path)
@@ -8,13 +8,13 @@ def size(path):
 
 
 class Condition(object):
-    format_string = ''
+    format_string: str = ''
 
-    def __init__(self, operator, *values):
+    def __init__(self, operator: str, *values) -> None:
         self.operator = operator
         self.values = values
 
-    def serialize(self, placeholder_names, expression_attribute_values):
+    def serialize(self, placeholder_names: Dict[str, str], expression_attribute_values: Dict[str, str]) -> str:
         values = [value.serialize(placeholder_names, expression_attribute_values) for value in self.values]
         return self.format_string.format(*values, operator=self.operator)
 
@@ -40,14 +40,10 @@ def __or__(self, other):
     def __invert__(self):
         return Not(self)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         values = [str(value) for value in self.values]
         return self.format_string.format(*values, operator=self.operator)
 
-    def __nonzero__(self):
-        # Prevent users from accidentally comparing the condition object instead of the attribute instance
-        raise TypeError("unsupported operand type(s) for bool: '{}'".format(self.__class__.__name__))
-
     def __bool__(self):
         # Prevent users from accidentally comparing the condition object instead of the attribute instance
         raise TypeError("unsupported operand type(s) for bool: {}".format(self.__class__.__name__))
@@ -59,19 +55,19 @@ class Comparison(Condition):
     def __init__(self, operator, lhs, rhs):
         if operator not in ['=', '<>', '<', '<=', '>', '>=']:
             raise ValueError("{0} is not a valid comparison operator: {0}".format(operator))
-        super(Comparison, self).__init__(operator, lhs, rhs)
+        super().__init__(operator, lhs, rhs)
 
 
 class Between(Condition):
     format_string = '{0} {operator} {1} AND {2}'
 
     def __init__(self, path, lower, upper):
-        super(Between, self).__init__(BETWEEN, path, lower, upper)
+        super().__init__('BETWEEN', path, lower, upper)
 
 
 class In(Condition):
     def __init__(self, path, *values):
-        super(In, self).__init__(IN, path, *values)
+        super().__init__('IN', path, *values)
         list_format = ', '.join('{' + str(i + 1) + '}' for i in range(len(values)))
         self.format_string = '{0} {operator} (' + list_format + ')'
 
@@ -80,53 +76,53 @@ class Exists(Condition):
     format_string = '{operator} ({0})'
 
     def __init__(self, path):
-        super(Exists, self).__init__('attribute_exists', path)
+        super().__init__('attribute_exists', path)
 
 
 class NotExists(Condition):
     format_string = '{operator} ({0})'
 
     def __init__(self, path):
-        super(NotExists, self).__init__('attribute_not_exists', path)
+        super().__init__('attribute_not_exists', path)
 
 
 class IsType(Condition):
     format_string = '{operator} ({0}, {1})'
 
     def __init__(self, path, attr_type):
-        super(IsType, self).__init__('attribute_type', path, attr_type)
+        super().__init__('attribute_type', path, attr_type)
 
 
 class BeginsWith(Condition):
     format_string = '{operator} ({0}, {1})'
 
     def __init__(self, path, prefix):
-        super(BeginsWith, self).__init__('begins_with', path, prefix)
+        super().__init__('begins_with', path, prefix)
 
 
 class Contains(Condition):
     format_string = '{operator} ({0}, {1})'
 
     def __init__(self, path, operand):
-        super(Contains, self).__init__('contains', path, operand)
+        super().__init__('contains', path, operand)
 
 
 class And(Condition):
     format_string = '({0} {operator} {1})'
 
-    def __init__(self, condition1, condition2):
-        super(And, self).__init__(AND, condition1, condition2)
+    def __init__(self, condition1: Condition, condition2: Condition) -> None:
+        super().__init__('AND', condition1, condition2)
 
 
 class Or(Condition):
     format_string = '({0} {operator} {1})'
 
-    def __init__(self, condition1, condition2):
-        super(Or, self).__init__(OR, condition1, condition2)
+    def __init__(self, condition1: Condition, condition2: Condition) -> None:
+        super().__init__('OR', condition1, condition2)
 
 
 class Not(Condition):
     format_string = '({operator} {0})'
 
-    def __init__(self, condition):
-        super(Not, self).__init__('NOT', condition)
+    def __init__(self, condition: Condition) -> None:
+        super().__init__('NOT', condition)

pynamodb/expressions/operand.py

@@ -15,15 +15,13 @@
 if TYPE_CHECKING:
     from pynamodb.attributes import Attribute
 
-_PathOrAttribute = Union['Path', 'Attribute', List[str], str]
-
 
 class _Operand:
     """
     Operand is the base class for objects that can be operands in Condition and Update Expressions.
     """
     format_string = ''
-    attr_type: Any = None
+    attr_type: Optional[str] = None
 
     def __init__(self, *values: Any) -> None:
         self.values = values
@@ -122,7 +120,7 @@ class _Size(_ConditionOperand):
     format_string = 'size ({0})'
     attr_type = NUMBER
 
-    def __init__(self, path: _PathOrAttribute) -> None:
+    def __init__(self, path: Union['Path', 'Attribute', str, List[str]]) -> None:
         if not isinstance(path, Path):
             path = Path(path)
         super(_Size, self).__init__(path)
@@ -235,9 +233,9 @@ class Path(_NumericOperand, _ListAppendOperand, _ConditionOperand):
     """
     format_string = '{0}'
 
-    def __init__(self, attribute_or_path: _PathOrAttribute) -> None:
+    def __init__(self, attribute_or_path: Union['Attribute', str, List[str]]) -> None:
         from pynamodb.attributes import Attribute  # prevent circular import -- Attribute imports Path
-        path: _PathOrAttribute
+        path: Union[str, List[str]]
         if isinstance(attribute_or_path, Attribute):
             self.attribute = attribute_or_path
             self.attr_type = attribute_or_path.attr_type
@@ -251,7 +249,7 @@ def __init__(self, attribute_or_path: _PathOrAttribute) -> None:
         super(Path, self).__init__(get_path_segments(path))
 
     @property
-    def path(self) -> Any:
+    def path(self) -> List[str]:
         return self.values[0]
 
     def __iter__(self):
@@ -338,6 +336,6 @@ def __repr__(self) -> str:
         return "Path({})".format(self.path)
 
     @staticmethod
-    def _quote_path(path):
+    def _quote_path(path: str) -> str:
         path, sep, rem = path.partition('[')
         return repr(path) + sep + rem

pynamodb/expressions/projection.py

@@ -1,16 +1,20 @@
+from typing import Dict
+from typing import List
+from typing import Union
+
 from pynamodb.attributes import Attribute
 from pynamodb.expressions.operand import Path
 from pynamodb.expressions.util import substitute_names
 
 
-def create_projection_expression(attributes_to_get, placeholders):
+def create_projection_expression(attributes_to_get, placeholders: Dict[str, str]) -> str:
     if not isinstance(attributes_to_get, list):
         attributes_to_get = [attributes_to_get]
     expressions = [substitute_names(_get_document_path(attribute), placeholders) for attribute in attributes_to_get]
     return ', '.join(expressions)
 
 
-def _get_document_path(attribute):
+def _get_document_path(attribute: Union[Attribute, Path, str]) -> List[str]:
     if isinstance(attribute, Attribute):
         return [attribute.attr_name]
     if isinstance(attribute, Path):

pynamodb/expressions/update.py

@@ -1,23 +1,31 @@
+from typing import Dict
 from typing import List
+from typing import Optional
+from typing import Sequence
 from typing import TYPE_CHECKING
 
-from pynamodb.constants import BINARY_SET, NUMBER, NUMBER_SET, STRING_SET
+from pynamodb.constants import BINARY_SET
+from pynamodb.constants import NUMBER
+from pynamodb.constants import NUMBER_SET
+from pynamodb.constants import STRING_SET
 
 if TYPE_CHECKING:
+    from pynamodb.expressions.operand import _Operand
     from pynamodb.expressions.operand import Path
+    from pynamodb.expressions.operand import Value
 
 
-class Action(object):
-    format_string = ''
+class Action:
+    format_string: str = ''
 
-    def __init__(self, *values: 'Path') -> None:
+    def __init__(self, *values: '_Operand') -> None:
         self.values = values
 
-    def serialize(self, placeholder_names, expression_attribute_values):
+    def serialize(self, placeholder_names: Dict[str, str], expression_attribute_values: Dict[str, str]) -> str:
         values = [value.serialize(placeholder_names, expression_attribute_values) for value in self.values]
         return self.format_string.format(*values)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         values = [str(value) for value in self.values]
         return self.format_string.format(*values)
 
@@ -28,7 +36,7 @@ class SetAction(Action):
     """
     format_string = '{0} = {1}'
 
-    def __init__(self, path: 'Path', value: 'Path') -> None:
+    def __init__(self, path: 'Path', value: '_Operand') -> None:
         super(SetAction, self).__init__(path, value)
 
 
@@ -48,7 +56,8 @@ class AddAction(Action):
     """
     format_string = '{0} {1}'
 
-    def __init__(self, path: 'Path', subset: 'Path') -> None:
+    def __init__(self, path: 'Path', subset: 'Value') -> None:
+        path._type_check(BINARY_SET, NUMBER, NUMBER_SET, STRING_SET)
         subset._type_check(BINARY_SET, NUMBER, NUMBER_SET, STRING_SET)
         super(AddAction, self).__init__(path, subset)
 
@@ -59,12 +68,13 @@ class DeleteAction(Action):
     """
     format_string = '{0} {1}'
 
-    def __init__(self, path: 'Path', subset: 'Path') -> None:
+    def __init__(self, path: 'Path', subset: 'Value') -> None:
+        path._type_check(BINARY_SET, NUMBER_SET, STRING_SET)
         subset._type_check(BINARY_SET, NUMBER_SET, STRING_SET)
         super(DeleteAction, self).__init__(path, subset)
 
 
-class Update(object):
+class Update:
 
     def __init__(self, *actions: Action) -> None:
         self.set_actions: List[SetAction] = []
@@ -86,22 +96,24 @@ def add_action(self, action: Action) -> None:
         else:
             raise ValueError("unsupported action type: '{}'".format(action.__class__.__name__))
 
-    def serialize(self, placeholder_names, expression_attribute_values):
-        expression = None
-        expression = self._add_clause(expression, 'SET', self.set_actions, placeholder_names, expression_attribute_values)
-        expression = self._add_clause(expression, 'REMOVE', self.remove_actions, placeholder_names, expression_attribute_values)
-        expression = self._add_clause(expression, 'ADD', self.add_actions, placeholder_names, expression_attribute_values)
-        expression = self._add_clause(expression, 'DELETE', self.delete_actions, placeholder_names, expression_attribute_values)
-        return expression
+    def serialize(self, placeholder_names: Dict[str, str], expression_attribute_values: Dict[str, str]) -> Optional[str]:
+        clauses = [
+            self._get_clause('SET', self.set_actions, placeholder_names, expression_attribute_values),
+            self._get_clause('REMOVE', self.remove_actions, placeholder_names, expression_attribute_values),
+            self._get_clause('ADD', self.add_actions, placeholder_names, expression_attribute_values),
+            self._get_clause('DELETE', self.delete_actions, placeholder_names, expression_attribute_values),
+        ]
+        expression = ' '.join(clause for clause in clauses if clause is not None)
+        return expression or None
 
     @staticmethod
-    def _add_clause(expression, keyword, actions, placeholder_names, expression_attribute_values):
-        clause = Update._get_clause(keyword, actions, placeholder_names, expression_attribute_values)
-        if clause is None:
-            return expression
-        return clause if expression is None else expression + " " + clause
-
-    @staticmethod
-    def _get_clause(keyword, actions, placeholder_names, expression_attribute_values):
-        actions = ", ".join([action.serialize(placeholder_names, expression_attribute_values) for action in actions])
-        return keyword + " " + actions if actions else None
+    def _get_clause(
+            keyword: str,
+            actions: Sequence[Action],
+            placeholder_names: Dict[str, str],
+            expression_attribute_values: Dict[str, str]
+    ) -> Optional[str]:
+        actions_string = ', '.join(
+            action.serialize(placeholder_names, expression_attribute_values) for action in actions
+        )
+        return keyword + ' ' + actions_string if actions_string else None

pynamodb/expressions/util.py

@@ -1,28 +1,36 @@
 import re
-from typing import Dict, List, Union
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from pynamodb.expressions.operands import Path
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Union
 
 
 PATH_SEGMENT_REGEX = re.compile(r'([^\[\]]+)((?:\[\d+\])*)$')
 
 
-def get_path_segments(document_path: Union[str, 'Path', List[str]]) -> Union[List[str], List['Path']]:
+def get_path_segments(document_path: Union[str, List[str]]) -> List[str]:
+    """
+    Splits a document path into nested elements using the map dereference operator (.)
+    and returns the list of path segments (an attribute name and optional list dereference operators ([n]).
+    If the document path is already a list of path segments, a new copy is returned.
+
+    https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.Attributes.html
+
+    Note: callers depend upon the returned list being a copy so that it may be safely mutated
+    """
     return document_path.split('.') if isinstance(document_path, str) else list(document_path)
 
 
-def substitute_names(document_path: Union[str, 'Path'], placeholders: Dict[str, str]) -> str:
+def substitute_names(document_path: Union[str, List[str]], placeholders: Dict[str, str]) -> str:
     """
     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 document_path: list of path segments (an attribute name and optional list dereference operators)
     :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,
+    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 = get_path_segments(document_path)
@@ -40,7 +48,7 @@ def substitute_names(document_path: Union[str, 'Path'], placeholders: Dict[str,
     return '.'.join(path_segments)
 
 
-def get_value_placeholder(value: 'Path', expression_attribute_values: Dict[str, str]) -> str:
+def get_value_placeholder(value: Any, expression_attribute_values: Dict[str, str]) -> str:
     placeholder = ':' + str(len(expression_attribute_values))
     expression_attribute_values[placeholder] = value
     return placeholder