Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion pynamodb_mypy/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
@@ -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
Expand Down
47 changes: 43 additions & 4 deletions tests/test_plugin.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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(
"""
Expand Down