models: range_key/filter_condition optional (pynamodb#588)

ikonst · web-flow · commit d19b1c163911 · 2019-01-17T02:18:05.000-08:00
The `range_key_condition` and `filter_condition` arguments default to `None`, and `None` triggers a specific behavior, so the type hints must allow passing `None`.
diff --git a/pynamodb/models.pyi b/pynamodb/models.pyi
@@ -58,8 +58,8 @@ class Model(metaclass=MetaModel):
     def query(
         cls: Type[_T],
         hash_key: KeyType,
-        range_key_condition: Condition = ...,
-        filter_condition: Condition = ...,
+        range_key_condition: Optional[Condition] = ...,
+        filter_condition: Optional[Condition] = ...,
         consistent_read: bool = ...,
         index_name: Optional[Text] = ...,
         scan_index_forward: Optional[Any] = ...,
diff --git a/pynamodb/tests/test_mypy.py b/pynamodb/tests/test_mypy.py
@@ -21,6 +21,30 @@ class MyModel(Model):
     """)
 
 
+def test_model_query():
+    assert_mypy_output("""
+    from pynamodb.attributes import NumberAttribute
+    from pynamodb.models import Model
+
+    class MyModel(Model):
+        my_attr = NumberAttribute()
+
+    # test hash key types
+    MyModel.query(123)
+    MyModel.query('123')
+    MyModel.query(12.3)
+    MyModel.query(b'123')
+    MyModel.query((1, 2, 3))
+    MyModel.query({'1': '2'})  # E: Argument 1 to "query" of "Model" has incompatible type "Dict[str, str]"; expected "Union[str, bytes, float, Tuple[Any, ...]]"
+
+    # test conditions
+    MyModel.query(123, range_key_condition=(MyModel.my_attr == 5), filter_condition=(MyModel.my_attr == 5))
+
+    # test conditions are optional
+    MyModel.query(123, range_key_condition=None, filter_condition=None)
+    """)  # noqa: E501
+
+
 def test_number_attribute():
     assert_mypy_output("""
     from pynamodb.attributes import NumberAttribute
