Fix type stubs for Model.count (pynamodb#587)

garrettheel · web-flow · commit 0d9b8e61a37e · 2019-01-16T14:56:59.000-08:00
diff --git a/pynamodb/models.py b/pynamodb/models.py
@@ -546,6 +546,7 @@ def count(cls,
         :param filter_condition: Condition used to restrict the query results
         :param consistent_read: If True, a consistent read is performed
         :param index_name: If set, then this index is used
+        :param rate_limit: If set then consumed capacity will be limited to this amount per second
         :param filters: A dictionary of filters to be used in the query. Requires a hash_key to be passed.
         """
         if hash_key is None:
diff --git a/pynamodb/models.pyi b/pynamodb/models.pyi
@@ -43,7 +43,17 @@ class Model(metaclass=MetaModel):
     @classmethod
     def from_raw_data(cls: Type[_T], data) -> _T: ...
     @classmethod
-    def count(cls: Type[_T], hash_key: Optional[KeyType] = ..., consistent_read: bool = ..., index_name: Optional[Text] = ..., limit: Optional[int] = ..., **filters) -> int: ...
+    def count(
+        cls: Type[_T],
+        hash_key: Optional[KeyType] = ...,
+        range_key_condition: Optional[Condition] = ...,
+        filter_condition: Optional[Condition] = ...,
+        consistent_read: bool = ...,
+        index_name: Optional[Text] = ...,
+        limit: Optional[int] = ...,
+        rate_limit: Optional[float] = ...,
+        **filters
+    ) -> int: ...
     @classmethod
     def query(
         cls: Type[_T],
diff --git a/pynamodb/tests/test_mypy.py b/pynamodb/tests/test_mypy.py
@@ -9,6 +9,18 @@
 from .mypy_helpers import assert_mypy_output  # noqa
 
 
+def test_model():
+    assert_mypy_output("""
+    from pynamodb.models import Model
+    from pynamodb.expressions.operand import Path
+
+    class MyModel(Model):
+        pass
+
+    reveal_type(MyModel.count('hash', Path('a').between(1, 3)))  # E: Revealed type is 'builtins.int'
+    """)
+
+
 def test_number_attribute():
     assert_mypy_output("""
     from pynamodb.attributes import NumberAttribute
