Skip to content

Commit

Permalink
Add type-checking to Manager.acreate (#2477)
Browse files Browse the repository at this point in the history
  • Loading branch information
vitosamson authored Jan 10, 2025
1 parent 93a6ef7 commit 5bff9f8
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 0 deletions.
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()

0 comments on commit 5bff9f8

Please sign in to comment.