diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 3c3084f..88201b4 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -4,7 +4,7 @@ from mypy.nodes import ( from mypy.plugin import ClassDefContext, DynamicClassDefContext from mypy.types import AnyType, Instance, TypeOfAny -from mypy_django_plugin.lib import helpers +from mypy_django_plugin.lib import fullnames, helpers def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None: @@ -65,9 +65,13 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte class_def_context = ClassDefContext(cls=new_manager_info.defn, reason=ctx.call, api=semanal_api) self_type = Instance(new_manager_info, []) - for name, sym in derived_queryset_info.names.items(): - if isinstance(sym.node, FuncDef): - helpers.copy_method_to_another_class(class_def_context, - self_type, - new_method_name=name, - method_node=sym.node) + # we need to copy all methods in MRO before django.db.models.query.QuerySet + for class_mro_info in derived_queryset_info.mro: + if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME: + break + for name, sym in class_mro_info.names.items(): + if isinstance(sym.node, FuncDef): + helpers.copy_method_to_another_class(class_def_context, + self_type, + new_method_name=name, + method_node=sym.node) diff --git a/test-data/typecheck/managers/querysets/test_from_queryset.yml b/test-data/typecheck/managers/querysets/test_from_queryset.yml index 689bfe3..8c339d6 100644 --- a/test-data/typecheck/managers/querysets/test_from_queryset.yml +++ b/test-data/typecheck/managers/querysets/test_from_queryset.yml @@ -94,3 +94,26 @@ class MyModel(models.Model): objects = NewManager() +- case: from_queryset_with_class_inheritance + main: | + from myapp.models import MyModel + reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_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 'builtins.str' + 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 + class BaseQuerySet(models.QuerySet): + def queryset_method(self) -> str: + return 'hello' + class ModelQuerySet(BaseQuerySet): + pass + + NewManager = BaseManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + objects = NewManager()