diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 4d362c1d33..c4b67e6fa4 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -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( diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml index 7d6e922710..5989ec4a12 100644 --- a/tests/typecheck/managers/querysets/test_as_manager.yml +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -1,3 +1,30 @@ +- 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_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 diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index ec08da3da1..f12ffb0c40 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -1,3 +1,25 @@ +- 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 + + 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