Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add type-checking to Manager.acreate #2477

Merged
merged 3 commits into from
Jan 10, 2025
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
1 change: 1 addition & 0 deletions mypy_django_plugin/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
20 changes: 20 additions & 0 deletions mypy_django_plugin/transformers/init_create.py
Original file line number Diff line number Diff line change
Expand Up @@ -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>]
model = get_proper_type(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")
18 changes: 18 additions & 0 deletions tests/typecheck/models/test_create.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
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
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()
Loading