Fix Index.query and Index.scan typing issues (#748)

Index.query and Index.scan should return a ResultIterator which iterates over instances of the respective model (which is not known at type-check type unless explicitly specified as a generic parameter).

Author: ikonst

docs/release_notes.rst

@@ -1,6 +1,13 @@
 Release Notes
 =============
 
+v4.3.1
+----------
+
+* Fix Index.query and Index.scan typing regressions introduced in 4.2.0, which were causing false errors
+  in type checkers
+
+
 v4.3.0
 ----------
 

pynamodb/__init__.py

@@ -7,4 +7,4 @@
 """
 __author__ = 'Jharrod LaFon'
 __license__ = 'MIT'
-__version__ = '4.3.0'
+__version__ = '4.3.1'

pynamodb/indexes.pyi

@@ -1,16 +1,17 @@
-from typing import Any, Dict, List, Optional, Text, Type, TypeVar
+from typing import Any, Dict, List, Optional, Text, TypeVar, Generic
 
 from pynamodb.expressions.condition import Condition
+from pynamodb.models import Model
 from pynamodb.pagination import ResultIterator
 
-_T = TypeVar('_T', bound='Index')
+_M = TypeVar('_M', bound=Model)
 
 
 class IndexMeta(type):
     def __init__(cls, name, bases, attrs) -> None: ...
 
 
-class Index(metaclass=IndexMeta):
+class Index(Generic[_M], metaclass=IndexMeta):
     Meta: Any
     def __init__(self) -> None: ...
     @classmethod
@@ -25,7 +26,7 @@ class Index(metaclass=IndexMeta):
     ) -> int: ...
     @classmethod
     def query(
-        cls: Type[_T],
+        cls,
         hash_key,
         range_key_condition: Optional[Condition] = ...,
         filter_condition: Optional[Condition] = ...,
@@ -36,10 +37,10 @@ class Index(metaclass=IndexMeta):
         attributes_to_get: Optional[Any] = ...,
         page_size: Optional[int] = ...,
         rate_limit: Optional[float] = ...,
-    ) -> ResultIterator[_T]: ...
+    ) -> ResultIterator[_M]: ...
     @classmethod
     def scan(
-        cls: Type[_T],
+        cls,
         filter_condition: Optional[Condition] = ...,
         segment: Optional[int] = ...,
         total_segments: Optional[int] = ...,
@@ -49,10 +50,10 @@ class Index(metaclass=IndexMeta):
         consistent_read: Optional[bool] = ...,
         rate_limit: Optional[float] = ...,
         attributes_to_get: Optional[List[str]] = ...,
-    ) -> ResultIterator[_T]: ...
+    ) -> ResultIterator[_M]: ...
 
-class GlobalSecondaryIndex(Index): ...
-class LocalSecondaryIndex(Index): ...
+class GlobalSecondaryIndex(Index[_M]): ...
+class LocalSecondaryIndex(Index[_M]): ...
 
 class Projection(object):
     projection_type: Any

pynamodb/models.pyi

@@ -1,4 +1,3 @@
-
 from .attributes import Attribute
 from .exceptions import DoesNotExist as DoesNotExist
 from typing import Any, Dict, Generic, Iterable, Iterator, List, Optional, Sequence, Tuple, Type, TypeVar, Text, Union

requirements-dev.txt

@@ -7,5 +7,5 @@ python-dateutil==2.8.0
 
 # only used in .travis.yml
 coveralls
-mypy==0.740;python_version>="3.7"
+mypy==0.761;python_version>="3.7"
 pytest-cov

tests/test_mypy.py

@@ -35,7 +35,7 @@ class MyModel(Model):
     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, ...]]"
+    MyModel.query({'1': '2'})  # E: Argument 1 to "query" of "Model" has incompatible type "Dict[str, str]"; expected "Union[str, bytes, float, int, Tuple[Any, ...]]"
 
     # test conditions
     MyModel.query(123, range_key_condition=(MyModel.my_attr == 5), filter_condition=(MyModel.my_attr == 5))
@@ -150,10 +150,52 @@ class MyModel(Model):
 
     reveal_type(MyModel.my_list)  # E: Revealed type is 'pynamodb.attributes.ListAttribute[__main__.MyMap]'
     reveal_type(MyModel().my_list)  # E: Revealed type is 'builtins.list[__main__.MyMap*]'
-    reveal_type(MyModel.my_list[0])  # E: Revealed type is 'Any'  # E: Value of type "ListAttribute[MyMap]" is not indexable
+    reveal_type(MyModel.my_list[0])  # E: Value of type "ListAttribute[MyMap]" is not indexable  # E: Revealed type is 'Any'
     reveal_type(MyModel().my_list[0].my_sub_attr)  # E: Revealed type is 'builtins.str'
 
     # Untyped lists are not well supported yet
-    reveal_type(MyModel.my_untyped_list[0])  # E: Revealed type is 'Any'  # E: Cannot determine type of 'my_untyped_list'
+    reveal_type(MyModel.my_untyped_list[0])  # E: Value of type "ListAttribute[Any]" is not indexable  # E: Revealed type is 'Any'
     reveal_type(MyModel().my_untyped_list[0].my_sub_attr)  # E: Revealed type is 'Any'
     """)
+
+
+def test_index_query_scan():
+    assert_mypy_output("""
+    from pynamodb.attributes import NumberAttribute
+    from pynamodb.models import Model
+    from pynamodb.indexes import GlobalSecondaryIndex
+    from pynamodb.pagination import ResultIterator
+
+    class UntypedIndex(GlobalSecondaryIndex):
+        bar = NumberAttribute(hash_key=True)
+
+    class TypedIndex(GlobalSecondaryIndex[MyModel]):
+        bar = NumberAttribute(hash_key=True)
+
+    class MyModel(Model):
+        foo = NumberAttribute(hash_key=True)
+        bar = NumberAttribute()
+
+        untyped_index = UntypedIndex()
+        typed_index = TypedIndex()
+
+    # Ensure old code keeps working
+    untyped_result: ResultIterator = MyModel.untyped_index.query(123)
+    model: MyModel = next(untyped_result)
+    not_model: int = next(untyped_result)  # this is legacy behavior so it's "fine"
+
+    # Allow users to specify which model their indices return
+    typed_result: ResultIterator[MyModel] = MyModel.typed_index.query(123)
+    my_model = next(typed_result)
+    not_model = next(typed_result)  # E: Incompatible types in assignment (expression has type "MyModel", variable has type "int")
+
+    # Ensure old code keeps working
+    untyped_result = MyModel.untyped_index.scan()
+    model = next(untyped_result)
+    not_model = next(untyped_result)  # this is legacy behavior so it's "fine"
+
+    # Allow users to specify which model their indices return
+    untyped_result = MyModel.typed_index.scan()
+    model = next(untyped_result)
+    not_model = next(untyped_result)  # E: Incompatible types in assignment (expression has type "MyModel", variable has type "int")
+    """)