Annotate Attribute.set typing (#1104)

Make sure that e.g. `MyModel.number_attr.set('spam')` fails type-check.

Author: conjmurph

pynamodb/attributes.py

@@ -253,7 +253,10 @@ def append(self, other: Iterable) -> '_ListAppend':
     def prepend(self, other: Iterable) -> '_ListAppend':
         return Path(self).prepend(other)
 
-    def set(self, value: Any) -> 'SetAction':
+    def set(
+        self,
+        value: Union[_T, 'Attribute[_T]', '_Increment', '_Decrement', '_IfNotExists', '_ListAppend']
+    ) -> 'SetAction':
         return Path(self).set(value)
 
     def remove(self) -> 'RemoveAction':

typing_tests/models.py

@@ -43,11 +43,12 @@ class MyModel(Model):
 
 
 def test_model_update() -> None:
-    from pynamodb.attributes import NumberAttribute
+    from pynamodb.attributes import NumberAttribute, UnicodeAttribute
     from pynamodb.models import Model
 
     class MyModel(Model):
         my_attr = NumberAttribute()
+        my_str_attr = UnicodeAttribute()
 
     my_model = MyModel()
     my_model.update(actions=[
@@ -59,6 +60,9 @@ class MyModel(Model):
         MyModel.my_attr.set(MyModel.my_attr | 123),
     ])
 
+    _ = MyModel.my_attr.set('foo')  # type:ignore[arg-type]
+    _ = MyModel.my_attr.set(MyModel.my_str_attr)  # type:ignore[arg-type]
+
 
 def test_paths() -> None:
     import pynamodb.expressions.operand