Support derived model initializers (#4)

Author: ikonst

pynamodb_mypy/plugin.py

@@ -103,10 +103,12 @@ def _get_function_signature_hook__pynamodb_model__init__(self, ctx: FunctionSigC
         if not isinstance(model_instance, mypy.types.Instance):  # pragma: no cover
             return ctx.default_signature
 
-        args = {
-            attr_name: _rehydrate_type(ctx.api, attr_data["type"])
-            for attr_name, attr_data in _pynamodb_attributes_metadata(model_instance.type).items()
-        }
+        args = {}
+        for model_cls in model_instance.type.mro:
+            args.update({
+                attr_name: _rehydrate_type(ctx.api, attr_data["type"])
+                for attr_name, attr_data in _pynamodb_attributes_metadata(model_cls).items()
+            })
 
         # substitute hash/range key types
         hash_key_type: mypy.types.Type = mypy.types.NoneTyp()

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = pynamodb-mypy
-version = 0.1.1
+version = 0.1.2
 description = mypy plugin for PynamoDB
 long_description = file: README.md
 long_description_content_type = text/markdown

tests/test_plugin.py

@@ -14,13 +14,21 @@ class MyModel(Model):
         my_range_key = NumberAttribute(range_key=True)
         my_attr = NumberAttribute()
 
+    class MyDerivedModel(MyModel):
+        my_derived_attr = NumberAttribute()
+
     MyModel(my_attr=5.5)
     MyModel(5.5, my_attr=5.5)
     MyModel(5.5, 5.5, my_attr=5.5)
     MyModel(hash_key=5.5, range_key=5.5, my_attr=5.5)
     MyModel(hash_key='hello', range_key='world', my_attr=5.5)  # E: Argument "hash_key" to "MyModel" has incompatible type "str"; expected "float"  [arg-type]
                                                                # E: Argument "range_key" to "MyModel" has incompatible type "str"; expected "float"  [arg-type]
     MyModel(foobar=5.5)  # E: Unexpected keyword argument "foobar" for "MyModel"  [call-arg]
+    
+    # test with derived model
+
+    MyDerivedModel(my_attr=5.5, my_derived_attr=42)
+    MyDerivedModel(foobar=5.5)  # E: Unexpected keyword argument "foobar" for "MyDerivedModel"  [call-arg]
     """
     )