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.

Sign up for GitHub

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

fill QuerySet generics using the manager's model type #2281

Merged
merged 2 commits into from
Jul 26, 2024
Merged

fill QuerySet generics using the manager's model type #2281

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9fa3404
fill QuerySet generics using the manager's model type
asottile Jul 26, 2024
658de86
construct as_manager instance at checker time
asottile Jul 26, 2024
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
16 changes: 7 additions & 9 deletions mypy_django_plugin/transformers/managers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
from mypy.semanal import SemanticAnalyzer
from mypy.semanal_shared import has_placeholder
from mypy.subtypes import find_member
from mypy.types import (
AnyType,
CallableType,
Expand All @@ -28,6 +29,7 @@
Overloaded,
ProperType,
TypeOfAny,
TypeType,
TypeVarType,
UnionType,
get_proper_type,
Expand Down Expand Up @@ -121,15 +123,11 @@ def _process_dynamic_method(
variables = method_type.variables
ret_type = method_type.ret_type

if not is_fallback_queryset:
queryset_instance = Instance(queryset_info, manager_instance.args)
else:
# The fallback queryset inherits _QuerySet, which has two generics
# instead of the one exposed on QuerySet. That means that we need
# to add the model twice. In real code it's not possible to inherit
# from _QuerySet, as it doesn't exist at runtime, so this fix is
# only needed for plugin-generated querysets.
queryset_instance = Instance(queryset_info, [manager_instance.args[0], manager_instance.args[0]])
manager_model = find_member("model", manager_instance, manager_instance)
assert isinstance(manager_model, TypeType), manager_model
manager_model_type = manager_model.item

queryset_instance = Instance(queryset_info, (manager_model_type,) * len(queryset_info.type_vars))

# For methods on the manager that return a queryset we need to override the
# return type to be the actual queryset class, not the base QuerySet that's
Expand Down
25 changes: 15 additions & 10 deletions tests/typecheck/managers/querysets/test_as_manager.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,19 +14,21 @@
- path: myapp/models.py
content: |
from django.db import models
from typing import List, Dict
from typing import List, Dict, TypeVar, ClassVar
from typing_extensions import Self

class BaseQuerySet(models.QuerySet):
M = TypeVar("M", bound=models.Model, covariant=True)

class BaseQuerySet(models.QuerySet[M]):
def example_dict(self) -> Dict[str, Self]: ...

class MyQuerySet(BaseQuerySet):
class MyQuerySet(BaseQuerySet[M]):
def example_simple(self) -> Self: ...
def example_list(self) -> List[Self]: ...
def just_int(self) -> int: ...

class MyModel(models.Model):
objects = MyQuerySet.as_manager()
objects = MyQuerySet.as_manager() # type: ignore[var-annotated]

class QuerySetWithoutSelf(models.QuerySet["MyModelWithoutSelf"]):
def method(self) -> "QuerySetWithoutSelf":
Expand Down Expand Up @@ -64,13 +66,16 @@
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from typing import TypeVar
from django.db import models

class MyQuerySet(models.QuerySet):
M = TypeVar("M", bound=models.Model, covariant=True)

class MyQuerySet(models.QuerySet[M]):
...

class MyModel(models.Model):
objects = MyQuerySet.as_manager()
objects = MyQuerySet.as_manager() # type: ignore[var-annotated]
asottile marked this conversation as resolved.
Show resolved Hide resolved

- case: model_gets_generated_manager_as_default_manager
main: |
Expand Down Expand Up @@ -183,7 +188,7 @@
from myapp.models import MyModel, MyModelManager
reveal_type(MyModelManager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[Any]"
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet"
installed_apps:
- myapp
files:
Expand All @@ -204,7 +209,7 @@
from myapp.models import MyModel, ManagerFromModelQuerySet
reveal_type(ManagerFromModelQuerySet) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[Any]"
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel]"
reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet"
installed_apps:
- myapp
files:
Expand Down Expand Up @@ -346,8 +351,8 @@
from myapp.models import MyModel
reveal_type(MyModel.objects_1) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects_2) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects_1.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects_2.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects_1.all()) # N: Revealed type is "myapp.models.ModelQuerySet"
reveal_type(MyModel.objects_2.all()) # N: Revealed type is "myapp.models.ModelQuerySet"
Copy link
Contributor Author

@asottile asottile Jul 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these (and all the similar ones) are more correct now -- these classes aren't generic so they shouldn't have type variables

I'm actually surprised mypy's internals don't enforce that an Instance has the correct number of typevars filled in!

installed_apps:
- myapp
files:
Expand Down
100 changes: 70 additions & 30 deletions tests/typecheck/managers/querysets/test_from_queryset.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,18 @@
content: |
from django.db import models
from django.db.models.manager import BaseManager
from typing import List, Dict
from typing import List, Dict, TypeVar
from typing_extensions import Self

class CustomManager(BaseManager):
M = TypeVar("M", covariant=True, bound=models.Model)

class CustomManager(BaseManager[M]):
def test_custom_manager(self) -> Self: ...

class BaseQuerySet(models.QuerySet):
class BaseQuerySet(models.QuerySet[M]):
def example_dict(self) -> Dict[str, Self]: ...

class MyQuerySet(BaseQuerySet):
class MyQuerySet(BaseQuerySet[M]):
def example_simple(self) -> Self: ...
def example_list(self) -> List[Self]: ...
def just_int(self) -> int: ...
Expand Down Expand Up @@ -82,10 +84,13 @@
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from typing import TypeVar
from django.db import models
from django.db.models.manager import BaseManager

class ModelQuerySet(models.QuerySet):
M = TypeVar("M", bound=models.Model, covariant=True)

class ModelQuerySet(models.QuerySet[M]):
def queryset_method(self) -> str:
return 'hello'
NewManager = BaseManager.from_queryset(ModelQuerySet)
Expand All @@ -103,7 +108,7 @@
reveal_type(MyModel.objects.queryset_method_3()) # N: Revealed type is "builtins.str"
reveal_type(MyModel.objects.queryset_method_4([])) # N: Revealed type is "None"
reveal_type(MyModel.objects.filter(id=1).queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet"
reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet"
installed_apps:
- myapp
files:
Expand Down Expand Up @@ -223,7 +228,7 @@
reveal_type(MyModel.objects.queryset_method_3()) # N: Revealed type is "builtins.str"
reveal_type(MyModel.objects.queryset_method_4([])) # N: Revealed type is "None"
reveal_type(MyModel.objects.filter(id=1).queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet"
reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet"
installed_apps:
- myapp
files:
Expand Down Expand Up @@ -307,7 +312,7 @@
import typing
kls: typing.Type[typing.Union[MyModel1, MyModel2]] = MyModel1
reveal_type(kls.objects) # N: Revealed type is "Union[myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel1], myapp.models.ManagerFromModelQuerySet2[myapp.models.MyModel2]]"
reveal_type(kls.objects.all()) # N: Revealed type is "Union[myapp.models.ModelQuerySet1[myapp.models.MyModel1], myapp.models.ModelQuerySet2[myapp.models.MyModel2]]"
reveal_type(kls.objects.all()) # N: Revealed type is "Union[myapp.models.ModelQuerySet1, myapp.models.ModelQuerySet2]"
reveal_type(kls.objects.get()) # N: Revealed type is "Union[myapp.models.MyModel1, myapp.models.MyModel2]"
reveal_type(kls.objects.queryset_method()) # N: Revealed type is "Union[builtins.int, builtins.str]"
installed_apps:
Expand Down Expand Up @@ -580,9 +585,9 @@
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.custom) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.all().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.custom().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
Expand Down Expand Up @@ -633,26 +638,26 @@
- case: from_queryset_includes_methods_returning_queryset
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects.alias) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.annotate) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.complex_filter) # N: Revealed type is "def (filter_obj: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.defer) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.difference) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.distinct) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.exclude) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.extra) # N: Revealed type is "def (select: Union[builtins.dict[builtins.str, Any], None] =, where: Union[typing.Sequence[builtins.str], None] =, params: Union[typing.Sequence[Any], None] =, tables: Union[typing.Sequence[builtins.str], None] =, order_by: Union[typing.Sequence[builtins.str], None] =, select_params: Union[typing.Sequence[Any], None] =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.intersection) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.none) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.only) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.order_by) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.prefetch_related) # N: Revealed type is "def (*lookups: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.reverse) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.select_for_update) # N: Revealed type is "def (nowait: builtins.bool =, skip_locked: builtins.bool =, of: typing.Sequence[builtins.str] =, no_key: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.select_related) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.union) # N: Revealed type is "def (*other_qs: Any, all: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.using) # N: Revealed type is "def (alias: Union[builtins.str, None]) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
reveal_type(MyModel.objects.alias) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.annotate) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.complex_filter) # N: Revealed type is "def (filter_obj: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.defer) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.difference) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.distinct) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.exclude) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.extra) # N: Revealed type is "def (select: Union[builtins.dict[builtins.str, Any], None] =, where: Union[typing.Sequence[builtins.str], None] =, params: Union[typing.Sequence[Any], None] =, tables: Union[typing.Sequence[builtins.str], None] =, order_by: Union[typing.Sequence[builtins.str], None] =, select_params: Union[typing.Sequence[Any], None] =) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.intersection) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.none) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.only) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.order_by) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.prefetch_related) # N: Revealed type is "def (*lookups: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.reverse) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.select_for_update) # N: Revealed type is "def (nowait: builtins.bool =, skip_locked: builtins.bool =, of: typing.Sequence[builtins.str] =, no_key: builtins.bool =) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.select_related) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.union) # N: Revealed type is "def (*other_qs: Any, all: builtins.bool =) -> myapp.models.MyQuerySet"
reveal_type(MyModel.objects.using) # N: Revealed type is "def (alias: Union[builtins.str, None]) -> myapp.models.MyQuerySet"
installed_apps:
- myapp
files:
Expand Down Expand Up @@ -882,6 +887,41 @@

class MCS(type): pass

- case: test_from_queryset_with_concrete_subclass
main: |
from myapp.models import Concrete
reveal_type(Concrete.objects) # N: Revealed type is "myapp.models.ConcreteManager"
reveal_type(Concrete.objects.get()) # N: Revealed type is "myapp.models.Concrete"
reveal_type(Concrete.objects.all()) # N: Revealed type is "myapp.models.CustomQuerySet[myapp.models.Concrete, myapp.models.Concrete]"
reveal_type(Concrete.objects.all().get()) # N: Revealed type is "myapp.models.Concrete"
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from typing import ClassVar
from typing_extensions import Self, TypeVar
from django.db.models import Model, QuerySet
from django.db.models.manager import Manager

M = TypeVar("M", bound=Model, covariant=True)
D = TypeVar("D", covariant=True, default=M)

class CustomQuerySet(QuerySet[M, D]): pass

_base = Manager.from_queryset(CustomQuerySet)

class CustomBase(_base[M]): ...

class BaseModel(Model):
objects: ClassVar[CustomBase[Self]] = CustomBase()

class ConcreteManager(CustomBase["Concrete"]): ...

class Concrete(BaseModel):
objects: ClassVar[ConcreteManager] = ConcreteManager()

- case: test_queryset_arg_as_unsupported_expressions
main: |
from typing import Union, Generic, TypeVar
Expand Down
Loading