Add inheritance QuerySet support for from_queryset (#275)

* Add testcase for queryset inheritance

* Add PoC

* Add condition for stop to looping over mro

* Change harcoded queryset class name to constant from fullnames

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():
+    # we need to copy all methods in MRO before django.db.models.query.QuerySet
-        if isinstance(sym.node, FuncDef):
+    for class_mro_info in derived_queryset_info.mro:
-            helpers.copy_method_to_another_class(class_def_context,
+        if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME:
-                                                 self_type,
+            break
-                                                 new_method_name=name,
+        for name, sym in class_mro_info.names.items():
-                                                 method_node=sym.node)
+            if isinstance(sym.node, FuncDef):
+                helpers.copy_method_to_another_class(class_def_context,
+                                                     self_type,
+                                                     new_method_name=name,
+                                                     method_node=sym.node)

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()