diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 17a227242..6b8e50276 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -243,26 +243,25 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: except helpers.IncompleteDefnException as exc: # Check if manager is a generated (dynamic class) manager base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0]) - if manager_fullname not in self.get_generated_manager_mappings(base_manager_fullname): + generated_managers = self.get_generated_manager_mappings(base_manager_fullname) + if manager_fullname not in generated_managers: # Manager doesn't appear to be generated. Track that we encountered an # incomplete definition and skip incomplete_manager_defs.add(manager_name) - continue + continue + + manager_info = self.lookup_typeinfo(generated_managers[manager_fullname]) + if manager_info is None: + continue - if manager_name not in self.model_classdef.info.names: + is_dynamically_generated = manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None + if manager_name not in self.model_classdef.info.names or is_dynamically_generated: manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, manager_type) - else: + elif self.has_any_parametrized_manager_as_base(manager_info): # Ending up here could for instance be due to having a custom _Manager_ # that is not built from a custom QuerySet. Another example is a # related manager. - # Don't interfere with dynamically generated manager classes - is_dynamically_generated = "django" in manager_info.metadata and manager_info.metadata["django"].get( - "from_queryset_manager" - ) - if not self.has_any_parametrized_manager_as_base(manager_info) or is_dynamically_generated: - continue - custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name try: custom_manager_type = self.create_new_model_parametrized_manager( diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 142f1f25a..2a2ae5db4 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -71,6 +71,55 @@ class MyModel(models.Model): objects = NewManager() +- case: from_queryset_generated_manager_imported_from_other_module + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects) # N: Revealed type is "myapp.querysets.NewManager[myapp.models.MyModel]" + reveal_type(MyModel.objects.get()) # N: Revealed type is "myapp.models.MyModel" + reveal_type(MyModel.objects.queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet" + reveal_type(MyModel.objects.queryset_method_2()) # N: Revealed type is "typing.Iterable[myapp.querysets.Custom]" + 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]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/querysets.py + content: | + from typing import TYPE_CHECKING, Iterable, Sequence + from django.db import models + from django.db.models.manager import BaseManager + if TYPE_CHECKING: + from .models import MyModel + + class Custom: + ... + + class ModelQuerySet(models.QuerySet["MyModel"]): + def queryset_method(self) -> "ModelQuerySet": + return self.filter() + + def queryset_method_2(self) -> Iterable[Custom]: + return [] + + def queryset_method_3(self) -> str: + return 'hello' + + def queryset_method_4(self, arg: Sequence) -> None: + return None + + NewManager = BaseManager.from_queryset(ModelQuerySet) + + - path: myapp/models.py + content: | + from django.db import models + from .querysets import NewManager + + class MyModel(models.Model): + objects = NewManager() + - case: from_queryset_with_manager main: | from myapp.models import MyModel