Type stubs for indexing into a ListAttribute in conditional expressions (#774)

Author: ikonst

docs/release_notes.rst

@@ -1,6 +1,19 @@
 Release Notes
 =============
 
+v4.3.3
+----------
+
+* Add type stubs for indexing into a ``ListAttribute`` for forming conditional expressions (#774)
+
+  ::
+
+    class MyModel(Model):
+      ...
+      my_list = ListAttribute()
+
+    MyModel.query(..., condition=MyModel.my_list[0] == 42)
+
 v4.3.2
 ----------
 

pynamodb/__init__.py

@@ -7,4 +7,4 @@
 """
 __author__ = 'Jharrod LaFon'
 __license__ = 'MIT'
-__version__ = '4.3.2'
+__version__ = '4.3.3'

pynamodb/attributes.pyi

@@ -6,7 +6,7 @@ from pynamodb.expressions.condition import (
     BeginsWith, Between, Comparison, Contains, NotExists, Exists, In
 )
 from pynamodb.expressions.operand import (
-    _Decrement, _IfNotExists, _Increment, _ListAppend
+    Path, _Decrement, _IfNotExists, _Increment, _ListAppend
 )
 from pynamodb.expressions.update import (
     AddAction, DeleteAction, RemoveAction, SetAction, ListRemoveAction
@@ -168,6 +168,7 @@ class ListAttribute(Generic[_T], Attribute[List[_T]]):
     def __get__(self: _A, instance: None, owner: Any) -> _A: ...
     @overload
     def __get__(self, instance: Any, owner: Any) -> List[_T]: ...
+    def __getitem__(self, idx: int) -> Path: ...
     def remove_indexes(self, *indexes: int) -> Union[ListRemoveAction]: ...
 
 

pynamodb/expressions/operand.pyi

@@ -79,3 +79,4 @@ class Path(_NumericOperand, _ListAppendOperand, _ConditionOperand):
     def set(self, value: Any) -> SetAction: ...
     def remove(self) -> RemoveAction: ...
     def remove_list_elements(self, *indexes: int) -> ListRemoveAction: ...
+    def __getitem__(self, idx: Union[int, str]) -> Path: ...

tests/test_mypy.py

@@ -142,15 +142,32 @@ 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: 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: 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_paths():
+    assert_mypy_output("""
+    from pynamodb.attributes import ListAttribute, MapAttribute, UnicodeAttribute
+    from pynamodb.models import Model
+
+    class MyMap(MapAttribute):
+        my_sub_attr = UnicodeAttribute()
+
+    class MyModel(Model):
+        my_list = ListAttribute(of=MyMap)
+        my_map = MyMap()
+
+    reveal_type(MyModel.my_list[0])  # E: Revealed type is 'pynamodb.expressions.operand.Path'
+    reveal_type(MyModel.my_list[0] == MyModel())  # E: Revealed type is 'pynamodb.expressions.condition.Comparison'
+    # the following string indexing is not type checked - not by mypy nor in runtime
+    reveal_type(MyModel.my_list[0]['my_sub_attr'] == 'foobar')  # E: Revealed type is 'pynamodb.expressions.condition.Comparison'
+    """)
+
+
 def test_index_query_scan():
     assert_mypy_output("""
     from pynamodb.attributes import NumberAttribute