From 00467f9826dbab82ecf236bfb485743fff0d31c4 Mon Sep 17 00:00:00 2001 From: Vito LaVilla Date: Fri, 10 Jan 2025 11:35:08 -0500 Subject: [PATCH 1/3] add type-checking to Manager.acreate --- mypy_django_plugin/main.py | 1 + .../transformers/init_create.py | 20 +++++++++++++++++++ tests/typecheck/models/test_create.yml | 18 +++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index b8290361b..088176cd8 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -172,6 +172,7 @@ def manager_and_queryset_method_hooks(self) -> dict[str, Callable[[MethodContext "alias": partial(querysets.extract_proper_type_queryset_annotate, django_context=self.django_context), "annotate": partial(querysets.extract_proper_type_queryset_annotate, django_context=self.django_context), "create": partial(init_create.redefine_and_typecheck_model_create, django_context=self.django_context), + "acreate": partial(init_create.redefine_and_typecheck_model_acreate, django_context=self.django_context), "filter": typecheck_filtering_method, "get": typecheck_filtering_method, "exclude": typecheck_filtering_method, diff --git a/mypy_django_plugin/transformers/init_create.py b/mypy_django_plugin/transformers/init_create.py index f26baaf6d..3c863afa9 100644 --- a/mypy_django_plugin/transformers/init_create.py +++ b/mypy_django_plugin/transformers/init_create.py @@ -76,3 +76,23 @@ def redefine_and_typecheck_model_create(ctx: MethodContext, django_context: Djan return ctx.default_return_type return typecheck_model_method(ctx, django_context, model_cls, "create") + + +def redefine_and_typecheck_model_acreate(ctx: MethodContext, django_context: DjangoContext) -> MypyType: + default_return_type = get_proper_type(ctx.default_return_type) + + if not isinstance(default_return_type, Instance): + # only work with ctx.default_return_type = model Instance + return ctx.default_return_type + + # default_return_type at this point should be of type Coroutine[Any, Any, ] + model = default_return_type.args[-1] + if not isinstance(model, Instance): + return ctx.default_return_type + + model_fullname = model.type.fullname + model_cls = django_context.get_model_class_by_fullname(model_fullname) + if model_cls is None: + return ctx.default_return_type + + return typecheck_model_method(ctx, django_context, model_cls, "acreate") diff --git a/tests/typecheck/models/test_create.yml b/tests/typecheck/models/test_create.yml index 91d28b1ab..c7b1fe0f7 100644 --- a/tests/typecheck/models/test_create.yml +++ b/tests/typecheck/models/test_create.yml @@ -155,3 +155,21 @@ id = models.IntegerField(primary_key=True) class MyModel3(models.Model): default = models.IntegerField(default=return_int) + +- case: default_manager_acreate_is_typechecked + main: | + import asyncio + from myapp.models import User + async def amain() -> None: + await User.objects.acreate(pk=1, name='Max', age=10) + await User.objects.acreate(age=[]) # E: Incompatible type for "age" of "User" (got "List[Any]", expected "Union[float, int, str, Combinable]") [misc] + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + class User(models.Model): + name = models.CharField(max_length=100) + age = models.IntegerField() From 201739cf2eb88548369f9abe60e7204b2eddc9cb Mon Sep 17 00:00:00 2001 From: Vito LaVilla Date: Fri, 10 Jan 2025 11:54:47 -0500 Subject: [PATCH 2/3] update test --- tests/typecheck/models/test_create.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/typecheck/models/test_create.yml b/tests/typecheck/models/test_create.yml index c7b1fe0f7..0192a5505 100644 --- a/tests/typecheck/models/test_create.yml +++ b/tests/typecheck/models/test_create.yml @@ -161,7 +161,7 @@ import asyncio from myapp.models import User async def amain() -> None: - await User.objects.acreate(pk=1, name='Max', age=10) + reveal_type(await User.objects.acreate(pk=1, name='Max', age=10)) # N: Revealed type is "myapp.models.User" await User.objects.acreate(age=[]) # E: Incompatible type for "age" of "User" (got "List[Any]", expected "Union[float, int, str, Combinable]") [misc] installed_apps: - myapp From 12da8029407e6fcdcab019120dae720f9f2b895b Mon Sep 17 00:00:00 2001 From: Vito LaVilla Date: Fri, 10 Jan 2025 11:56:04 -0500 Subject: [PATCH 3/3] use get_proper_type --- mypy_django_plugin/transformers/init_create.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypy_django_plugin/transformers/init_create.py b/mypy_django_plugin/transformers/init_create.py index 3c863afa9..19647d261 100644 --- a/mypy_django_plugin/transformers/init_create.py +++ b/mypy_django_plugin/transformers/init_create.py @@ -86,7 +86,7 @@ def redefine_and_typecheck_model_acreate(ctx: MethodContext, django_context: Dja return ctx.default_return_type # default_return_type at this point should be of type Coroutine[Any, Any, ] - model = default_return_type.args[-1] + model = get_proper_type(default_return_type.args[-1]) if not isinstance(model, Instance): return ctx.default_return_type