Skip to content

Commit

Permalink
fix as manager and from queryset self types (#1788) (#1789)
Browse files Browse the repository at this point in the history
moranabadie authored Oct 21, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
1 parent 1334efd commit 0ce571e
Showing 3 changed files with 54 additions and 0 deletions.
3 changes: 3 additions & 0 deletions mypy_django_plugin/transformers/managers.py
Original file line number Diff line number Diff line change
@@ -132,6 +132,9 @@ def _process_dynamic_method(
_replace_type_var(arg_type, base_that_has_method.defn.type_vars[0].fullname, manager_instance.args[0])
for arg_type in args_types
]
if base_that_has_method.self_type:
# Manages -> Self returns
ret_type = _replace_type_var(ret_type, base_that_has_method.self_type.fullname, manager_instance)

# Drop any 'self' argument as our manager is already initialized
return method_type.copy_modified(
28 changes: 28 additions & 0 deletions tests/typecheck/managers/querysets/test_as_manager.yml
Original file line number Diff line number Diff line change
@@ -1,3 +1,31 @@
- case: self_return_management
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects.example_simple()) # N: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.example_list()) # N: Revealed type is "builtins.list[myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]]"
reveal_type(MyModel.objects.example_simple().just_int()) # N: Revealed type is "builtins.int"
reveal_type(MyModel.objects.example_dict()) # N: Revealed type is "builtins.dict[builtins.str, myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]]"
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from typing import List, Dict
from typing_extensions import Self
class BaseQuerySet(models.QuerySet):
def example_dict(self) -> Dict[str, Self]: ...
class MyQuerySet(BaseQuerySet):
def example_simple(self) -> Self: ...
def example_list(self) -> List[Self]: ...
def just_int(self) -> int: ...
class MyModel(models.Model):
objects = MyQuerySet.as_manager()
- case: declares_manager_type_like_django
main: |
from myapp.models import MyModel
23 changes: 23 additions & 0 deletions tests/typecheck/managers/querysets/test_from_queryset.yml
Original file line number Diff line number Diff line change
@@ -1,3 +1,26 @@
- case: from_queryset_self_return_management
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects.example_simple()) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.example_list()) # N: Revealed type is "builtins.list[myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]]"
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from django.db.models.manager import BaseManager
from typing_extensions import Self
from typing import List
class ModelQuerySet(models.QuerySet):
def example_simple(self) -> Self: ...
def example_list(self) -> List[Self]: ...
NewManager = BaseManager.from_queryset(ModelQuerySet)
class MyModel(models.Model):
objects = NewManager()
- case: from_queryset_with_base_manager
main: |
from myapp.models import MyModel

0 comments on commit 0ce571e

Please sign in to comment.