diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index ce41b6ee9..7999354f2 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -87,11 +87,21 @@ def get_funcdef_type(definition: Union[FuncBase, Decorator, None]) -> Optional[P variables = method_type.variables ret_type = method_type.ret_type + is_fallback_queryset = queryset_info.metadata.get("django", {}).get("any_fallback_queryset", False) + # 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 # used by the typing stubs. if method_name in MANAGER_METHODS_RETURNING_QUERYSET: - ret_type = Instance(queryset_info, manager_instance.args) + if not is_fallback_queryset: + ret_type = 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 pluign-generated querysets. + ret_type = Instance(queryset_info, [manager_instance.args[0], manager_instance.args[0]]) variables = [] # Drop any 'self' argument as our manager is already initialized diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index a4d2f7a1d..9c468d32c 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -12,6 +12,7 @@ from mypy.types import AnyType, Instance from mypy.types import Type as MypyType from mypy.types import TypedDictType, TypeOfAny +from mypy.typevars import fill_typevars from mypy_django_plugin.django.context import DjangoContext from mypy_django_plugin.errorcodes import MANAGER_MISSING @@ -19,7 +20,10 @@ from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME, MODEL_CLASS_FULLNAME from mypy_django_plugin.transformers import fields from mypy_django_plugin.transformers.fields import get_field_descriptor_types -from mypy_django_plugin.transformers.managers import create_manager_info_from_from_queryset_call +from mypy_django_plugin.transformers.managers import ( + MANAGER_METHODS_RETURNING_QUERYSET, + create_manager_info_from_from_queryset_call, +) class ModelClassInitializer: @@ -83,6 +87,85 @@ def get_generated_manager_info(self, manager_fullname: str, base_manager_fullnam # Not a generated manager return None + def get_or_create_manager_with_any_fallback(self, related_manager: bool = False) -> TypeInfo: + """ + Create a Manager subclass with fallback to Any for unknown attributes + and methods. This is used for unresolved managers, where we don't know + the actual type of the manager. + + The created class is reused if multiple unknown managers are encountered. + """ + + name = "UnknownManager" if not related_manager else "UnknownRelatedManager" + + # Check if we've already created a fallback manager class for this + # module, and if so reuse that. + manager_info = self.lookup_typeinfo(f"{self.model_classdef.info.module_name}.{name}") + if manager_info and manager_info.metadata.get("django", {}).get("any_fallback_manager"): + return manager_info + + fallback_queryset = self.get_or_create_queryset_with_any_fallback() + base_manager_fullname = ( + fullnames.MANAGER_CLASS_FULLNAME if not related_manager else fullnames.RELATED_MANAGER_CLASS + ) + base_manager_info = self.lookup_typeinfo(base_manager_fullname) + assert base_manager_info, f"Type info for {base_manager_fullname} missing" + + base_manager = fill_typevars(base_manager_info) + assert isinstance(base_manager, Instance) + manager_info = self.add_new_class_for_current_module(name, [base_manager]) + manager_info.fallback_to_any = True + + manager_info.type_vars = base_manager_info.type_vars + manager_info.defn.type_vars = base_manager_info.defn.type_vars + manager_info.metaclass_type = manager_info.calculate_metaclass_type() + + # For methods on BaseManager that return a queryset we need to update + # the return type to be the actual queryset subclass used. This is done + # by adding the methods as attributes with type Any to the manager + # class. The actual type of these methods are resolved in + # resolve_manager_method. + for method_name in MANAGER_METHODS_RETURNING_QUERYSET: + helpers.add_new_sym_for_info(manager_info, name=method_name, sym_type=AnyType(TypeOfAny.special_form)) + + manager_info.metadata["django"] = django_metadata = { + "any_fallback_manager": True, + "from_queryset_manager": fallback_queryset.fullname, + } + + return manager_info + + def get_or_create_queryset_with_any_fallback(self) -> TypeInfo: + """ + Create a QuerySet subclass with fallback to Any for unknown attributes + and methods. This is used for the manager returned by the method above. + """ + + name = "UnknownQuerySet" + + # Check if we've already created a fallback queryset class for this + # module, and if so reuse that. + queryset_info = self.lookup_typeinfo(f"{self.model_classdef.info.module_name}.{name}") + if queryset_info and queryset_info.metadata.get("django", {}).get("any_fallback_queryset"): + return queryset_info + + base_queryset_info = self.lookup_typeinfo(fullnames.QUERYSET_CLASS_FULLNAME) + assert base_queryset_info, f"Type info for {fullnames.QUERYSET_CLASS_FULLNAME} missing" + + base_queryset = fill_typevars(base_queryset_info) + assert isinstance(base_queryset, Instance) + queryset_info = self.add_new_class_for_current_module(name, [base_queryset]) + queryset_info.metadata["django"] = { + "any_fallback_queryset": True, + } + queryset_info.fallback_to_any = True + + queryset_info.type_vars = base_queryset_info.type_vars.copy() + queryset_info.defn.type_vars = base_queryset_info.defn.type_vars.copy() + queryset_info.metaclass_type = queryset_info.calculate_metaclass_type() + + return queryset_info + def run_with_model_cls(self, model_cls): raise NotImplementedError("Implement this in subclasses") @@ -285,13 +368,11 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: # (Django's default) before finishing. And emit an error, to allow for # ignoring a more specialised manager not being resolved while still # setting _some_ type - django_manager_info = self.lookup_typeinfo(fullnames.MANAGER_CLASS_FULLNAME) - assert ( - django_manager_info is not None - ), f"Type info for Django's {fullnames.MANAGER_CLASS_FULLNAME} missing" + fallback_manager_info = self.get_or_create_manager_with_any_fallback() self.add_new_node_to_model_class( - manager_name, Instance(django_manager_info, [Instance(self.model_classdef.info, [])]) + manager_name, Instance(fallback_manager_info, [Instance(self.model_classdef.info, [])]) ) + # Find expression for e.g. `objects = SomeManager()` manager_expr = self.get_manager_expression(manager_name) manager_fullname = f"{self.model_classdef.fullname}.{manager_name}" @@ -425,28 +506,27 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: except helpers.IncompleteDefnException as exc: if not self.api.final_iteration: raise exc - else: - if related_manager_info: - """ - If a django model has a Manager class that cannot be - resolved statically (if it is generated in a way - where we cannot import it, like `objects = my_manager_factory()`), - we fallback to the default related manager, so you - at least get a base level of working type checking. - - See https://github.com/typeddjango/django-stubs/pull/993 - for more information on when this error can occur. - """ - self.add_new_node_to_model_class( - attname, Instance(related_manager_info, [Instance(related_model_info, [])]) - ) - related_model_fullname = related_model_cls.__module__ + "." + related_model_cls.__name__ - self.ctx.api.fail( - f"Couldn't resolve related manager for relation {relation.name!r} (from {related_model_fullname}.{relation.field}).", - self.ctx.cls, - code=MANAGER_MISSING, - ) - continue + + # If a django model has a Manager class that cannot be + # resolved statically (if it is generated in a way where we + # cannot import it, like `objects = my_manager_factory()`), + # we fallback to the default related manager, so you at + # least get a base level of working type checking. + # + # See https://github.com/typeddjango/django-stubs/pull/993 + # for more information on when this error can occur. + fallback_manager = self.get_or_create_manager_with_any_fallback(related_manager=True) + self.add_new_node_to_model_class( + attname, Instance(fallback_manager, [Instance(related_model_info, [])]) + ) + related_model_fullname = related_model_cls.__module__ + "." + related_model_cls.__name__ + self.ctx.api.fail( + f"Couldn't resolve related manager for relation {relation.name!r} (from {related_model_fullname}.{relation.field}).", + self.ctx.cls, + code=MANAGER_MISSING, + ) + + continue # Check if the related model has a related manager subclassed from the default manager # TODO: Support other reverse managers than `_default_manager` diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index c7f726732..50410d52e 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -678,18 +678,24 @@ def custom(self) -> None: pass + def TransactionManager(): + return BaseManager.from_queryset(TransactionQuerySet)() + class Transaction(models.Model): - objects = BaseManager.from_queryset(TransactionQuerySet) + objects = TransactionManager() def test(self) -> None: reveal_type(self.transactionlog_set) # We use a fallback Any type: + reveal_type(Transaction.objects) reveal_type(Transaction.objects.custom()) class TransactionLog(models.Model): transaction = models.ForeignKey(Transaction, on_delete=models.CASCADE) out: | - myapp/models:10: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]" - myapp/models:12: note: Revealed type is "Any" + myapp/models:11: error: Could not resolve manager type for "myapp.models.Transaction.objects" + myapp/models:13: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]" + myapp/models:15: note: Revealed type is "myapp.models.UnknownManager[myapp.models.Transaction]" + myapp/models:16: note: Revealed type is "Any" - case: resolve_primary_keys_for_foreign_keys_with_abstract_self_model diff --git a/tests/typecheck/managers/test_managers.yml b/tests/typecheck/managers/test_managers.yml index 51ee4467e..f79cfb283 100644 --- a/tests/typecheck/managers/test_managers.yml +++ b/tests/typecheck/managers/test_managers.yml @@ -502,6 +502,24 @@ reveal_type(user.bookingowner_set) reveal_type(user.booking_set) + + # Check QuerySet methods on UnknownManager + reveal_type(Booking.objects.all) + reveal_type(Booking.objects.custom) + reveal_type(Booking.objects.all().filter) + reveal_type(Booking.objects.all().custom) + reveal_type(Booking.objects.first()) + reveal_type(Booking.objects.get()) + reveal_type([booking for booking in Booking.objects.all()]) + reveal_type([booking for booking in Booking.objects.all().filter()]) + + + # Check QuerySet methods on UnknownRelatedManager + reveal_type(user.booking_set.all) + reveal_type(user.booking_set.custom) + reveal_type(user.booking_set.all().filter) + reveal_type(user.booking_set.all().custom) + reveal_type(user.booking_set.all().first()) out: | myapp/models:13: error: Couldn't resolve related manager for relation 'booking' (from myapp.models.Booking.myapp.Booking.renter). myapp/models:13: error: Couldn't resolve related manager for relation 'bookingowner_set' (from myapp.models.Booking.myapp.Booking.owner). @@ -512,12 +530,25 @@ myapp/models:32: error: Could not resolve manager type for "myapp.models.InvisibleUnresolvable.objects" myapp/models:36: note: Revealed type is "django.db.models.manager.Manager[myapp.models.User]" myapp/models:37: note: Revealed type is "django.db.models.manager.Manager[myapp.models.User]" - myapp/models:39: note: Revealed type is "django.db.models.manager.Manager[myapp.models.Booking]" + myapp/models:39: note: Revealed type is "myapp.models.UnknownManager[myapp.models.Booking]" myapp/models:40: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.Booking]" - myapp/models:42: note: Revealed type is "django.db.models.manager.Manager[myapp.models.TwoUnresolvable]" - myapp/models:43: note: Revealed type is "django.db.models.manager.Manager[myapp.models.TwoUnresolvable]" + myapp/models:42: note: Revealed type is "myapp.models.UnknownManager[myapp.models.TwoUnresolvable]" + myapp/models:43: note: Revealed type is "myapp.models.UnknownManager[myapp.models.TwoUnresolvable]" myapp/models:44: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.TwoUnresolvable]" - myapp/models:46: note: Revealed type is "django.db.models.manager.Manager[myapp.models.InvisibleUnresolvable]" + myapp/models:46: note: Revealed type is "myapp.models.UnknownManager[myapp.models.InvisibleUnresolvable]" myapp/models:47: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.InvisibleUnresolvable]" - myapp/models:49: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]" - myapp/models:50: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]" + myapp/models:49: note: Revealed type is "myapp.models.UnknownRelatedManager[myapp.models.Booking]" + myapp/models:50: note: Revealed type is "myapp.models.UnknownRelatedManager[myapp.models.Booking]" + myapp/models:53: note: Revealed type is "def () -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]" + myapp/models:54: note: Revealed type is "Any" + myapp/models:55: note: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]" + myapp/models:56: note: Revealed type is "Any" + myapp/models:57: note: Revealed type is "Union[myapp.models.Booking, None]" + myapp/models:58: note: Revealed type is "myapp.models.Booking" + myapp/models:59: note: Revealed type is "builtins.list[myapp.models.Booking]" + myapp/models:60: note: Revealed type is "builtins.list[myapp.models.Booking]" + myapp/models:64: note: Revealed type is "def () -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]" + myapp/models:65: note: Revealed type is "Any" + myapp/models:66: note: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]" + myapp/models:67: note: Revealed type is "Any" + myapp/models:68: note: Revealed type is "Union[myapp.models.Booking, None]"