Handle unexpected Model.__init__ sig (#3)

Author: ikonst

pynamodb_mypy/plugin.py

@@ -122,7 +122,9 @@ def _get_function_signature_hook__pynamodb_model__init__(self, ctx: FunctionSigC
             hash_key_idx = ctx.default_signature.arg_names.index("hash_key")
             range_key_idx = ctx.default_signature.arg_names.index("range_key")
             kwargs_idx = ctx.default_signature.arg_kinds.index(ArgKind.ARG_STAR2)
-        except IndexError:  # pragma: no cover
+        except ValueError:
+            ctx.api.fail(f"Unexpected signature '{ctx.default_signature}' for a PynamoDB model initializer: "
+                         "expecting 'hash_key', 'range_key' and a keywords argument", ctx.context)
             return ctx.default_signature
         else:
             arg_kinds = ctx.default_signature.arg_kinds.copy()

setup.cfg

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

tests/test_plugin.py

@@ -1,11 +1,9 @@
 from __future__ import annotations
 
-from unittest.mock import Mock
-
 from .mypy_helpers import MypyAssert
 
 
-def test_init(assert_mypy_output: MypyAssert) -> None:
+def test_model_init(assert_mypy_output: MypyAssert) -> None:
     assert_mypy_output(
         """
     from pynamodb.attributes import NumberAttribute
@@ -27,7 +25,7 @@ class MyModel(Model):
     )
 
 
-def test_init__no_attributes(assert_mypy_output: MypyAssert) -> None:
+def test_model_init__no_attributes(assert_mypy_output: MypyAssert) -> None:
     assert_mypy_output(
         """
     from pynamodb.attributes import NumberAttribute
@@ -45,6 +43,47 @@ class MyModel(Model):
     )
 
 
+def test_model_init__custom_empty(assert_mypy_output: MypyAssert) -> None:
+    assert_mypy_output(
+        """
+    from pynamodb.attributes import NumberAttribute
+    from pynamodb.models import Model
+
+    class MyModel(Model):
+        my_hash_key = NumberAttribute(hash_key=True)
+        my_range_key = NumberAttribute(range_key=True)
+        my_attr = NumberAttribute()
+
+        def __init__(self) -> None:
+          ...
+
+    MyModel('foo', 'bar')  # E: Unexpected signature 'def () -> __main__.MyModel' for a PynamoDB model initializer: expecting 'hash_key', 'range_key' and a keywords argument  [misc]
+                           # E: Too many arguments for "MyModel"  [call-arg]
+    """
+    )
+
+
+def test_model_init__custom_all_args(assert_mypy_output: MypyAssert) -> None:
+    assert_mypy_output(
+        """
+    from typing import Any
+
+    from pynamodb.attributes import NumberAttribute
+    from pynamodb.models import Model
+
+    class MyModel(Model):
+        my_hash_key = NumberAttribute(hash_key=True)
+        my_range_key = NumberAttribute(range_key=True)
+        my_attr = NumberAttribute()
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+          ...
+
+    MyModel(unknown=42)  # E: Unexpected signature 'def (*args: Any, **kwargs: Any) -> __main__.MyModel' for a PynamoDB model initializer: expecting 'hash_key', 'range_key' and a keywords argument  [misc]
+    """
+    )
+
+
 def test_number_attribute(assert_mypy_output: MypyAssert) -> None:
     assert_mypy_output(
         """