diff --git a/pynamodb_mypy/plugin.py b/pynamodb_mypy/plugin.py index b1d2a54..77798c2 100644 --- a/pynamodb_mypy/plugin.py +++ b/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() diff --git a/setup.cfg b/setup.cfg index d6d1b00..89380f9 100644 --- a/setup.cfg +++ b/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 diff --git a/tests/test_plugin.py b/tests/test_plugin.py index be25195..6ddfe23 100644 --- a/tests/test_plugin.py +++ b/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( """