Update to newer django-stubs and remove hack

extended_mypy_django_plugin/plugin/_plugin.py

@@ -1,4 +1,3 @@
-import collections
 import sys
 from typing import Generic
 
@@ -9,7 +8,6 @@
 from mypy.plugin import (
     AnalyzeTypeContext,
     AttributeContext,
-    ClassDefContext,
     DynamicClassDefContext,
     FunctionContext,
     FunctionSigContext,
@@ -28,16 +26,7 @@
 )
 from typing_extensions import assert_never
 
-from . import (
-    _config,
-    _dependencies,
-    _helpers,
-    _hook,
-    _known_annotations,
-    _reports,
-    _store,
-    actions,
-)
+from . import _config, _dependencies, _hook, _known_annotations, _reports, _store, actions
 
 
 class Hook(
@@ -151,65 +140,6 @@ def get_additional_deps(self, file: MypyFile) -> list[tuple[int, str, int]]:
         )
         return results
 
-    @_hook.hook
-    class get_base_class_hook(Hook[ClassDefContext, None]):
-        """
-        We need to make up for a bug in django-stubs
-        """
-
-        def choose(self) -> bool:
-            if self.super_hook is None:
-                return False
-
-            if _helpers.get_is_abstract_model() is None:
-                return False
-
-            sym = self.plugin.lookup_fully_qualified(self.fullname)
-            return bool(
-                sym is not None
-                and isinstance(sym.node, TypeInfo)
-                and _helpers.is_model_type(sym.node)
-            )
-
-        def run(self, ctx: ClassDefContext) -> None:
-            if self.super_hook is None:
-                return None
-
-            # Copy the code in django-stubs that crashes
-            # And fill in the missing information before continuing
-            processed_models = set()
-            model_bases = collections.deque([ctx.cls])
-            while model_bases:
-                model = model_bases.popleft()
-
-                try:
-                    # Whether this causes an AssertionError or an AttributeError depends
-                    # on whether mypy is compiled or not
-                    # Note that this only appears to trigger on followup changes with a cache
-                    # in very specific situations
-                    for base in model.info.bases:
-                        break
-                except AssertionError as exc:
-                    if str(exc) == "ClassDef is lacking info":
-                        sym = self.plugin.lookup_fully_qualified(model.fullname)
-                        if sym and isinstance(sym.node, TypeInfo):
-                            model.info = sym.node
-                except AttributeError as exc:
-                    if str(exc) == "attribute 'bases' of 'TypeInfo' undefined":
-                        sym = self.plugin.lookup_fully_qualified(model.fullname)
-                        if sym and isinstance(sym.node, TypeInfo):
-                            model.info = sym.node
-
-                for base in model.info.bases:
-                    if (
-                        _helpers.is_abstract_model(base.type)
-                        and base.type.fullname not in processed_models
-                    ):
-                        model_bases.append(base.type.defn)
-                        processed_models.add(base.type.fullname)
-
-            return self.super_hook(ctx)
-
     @_hook.hook
     class get_dynamic_class_hook(Hook[DynamicClassDefContext, None]):
         """

extended_mypy_django_plugin/plugin/_store.py

@@ -1,3 +1,4 @@
+import importlib.metadata
 from collections.abc import Iterator, Mapping, Sequence
 from typing import Protocol
 
@@ -8,9 +9,12 @@
 
 from ._reports import ModelModules
 
-QUERYSET_CLASS_FULLNAME = "django.db.models.query._QuerySet"
 MODEL_CLASS_FULLNAME = "django.db.models.base.Model"
 
+QUERYSET_CLASS_FULLNAME = "django.db.models.query.QuerySet"
+if importlib.metadata.version("mypy") == "1.4.0":
+    QUERYSET_CLASS_FULLNAME = "django.db.models.query._QuerySet"
+
 
 class UnionMustBeOfTypes(Exception):
     pass
@@ -99,7 +103,7 @@ def realise_querysets(
         for fullname, model in querysets:
             queryset = lookup_info(fullname)
             if not queryset:
-                raise RestartDmypy()
+                raise RestartDmypy(f"Could not find queryset for {fullname}")
 
             if not queryset.is_generic():
                 yield Instance(queryset, [])
@@ -184,7 +188,7 @@ def _get_dynamic_manager(
         """
         model_cls = self._get_model_class_by_fullname(model.fullname)
         if model_cls is None:
-            raise RestartDmypy()
+            raise RestartDmypy(f"Could not find model class for {model.fullname}")
 
         manager = model_cls._default_manager
         if manager is None:
@@ -203,7 +207,7 @@ def _get_dynamic_manager(
 
             base_manager_info = lookup_info(base_manager_fullname)
             if not base_manager_info:
-                raise RestartDmypy()
+                raise RestartDmypy(f"Could not find base manager for {base_manager_fullname}")
 
             metadata = base_manager_info.metadata
 

extended_mypy_django_plugin/plugin/actions/_type_analyze.py

@@ -67,8 +67,8 @@ def find_concrete_querysets(self, unbound_type: UnboundType) -> MypyType:
 
         try:
             querysets = tuple(self.store.realise_querysets(UnionType(concrete), self.lookup_info))
-        except _store.RestartDmypy:
-            self.api.fail("You probably need to restart dmypy", unbound_type)
+        except _store.RestartDmypy as err:
+            self.api.fail(f"You probably need to restart dmypy: {err}", unbound_type)
             return AnyType(TypeOfAny.from_error)
         except _store.UnionMustBeOfTypes:
             self.api.fail("Union must be of instances of models", unbound_type)
@@ -93,8 +93,8 @@ def find_default_queryset(self, unbound_type: UnboundType) -> MypyType:
 
         try:
             querysets = tuple(self.store.realise_querysets(type_arg, self.lookup_info))
-        except _store.RestartDmypy:
-            self.api.fail("You probably need to restart dmypy", unbound_type)
+        except _store.RestartDmypy as err:
+            self.api.fail(f"You probably need to restart dmypy: {err}", unbound_type)
             return AnyType(TypeOfAny.from_error)
         except _store.UnionMustBeOfTypes:
             self.api.fail("Union must be of instances of models", unbound_type)

extended_mypy_django_plugin/plugin/actions/_type_checker.py

@@ -312,8 +312,8 @@ def get_concrete_queryset_return_type(
             assert isinstance(concrete, Instance)
             try:
                 result.extend(self.store.realise_querysets(concrete, self.lookup_info))
-            except _store.RestartDmypy:
-                self.api.fail("You probably need to restart dmypy", context)
+            except _store.RestartDmypy as err:
+                self.api.fail(f"You probably need to restart dmypy: {err}", context)
                 return AnyType(TypeOfAny.from_error)
             except _store.UnionMustBeOfTypes:
                 self.api.fail("Union must be of instances of models", context)
@@ -326,8 +326,8 @@ def get_default_queryset_return_type(
     ) -> MypyType:
         try:
             querysets = tuple(self.store.realise_querysets(instances, self.lookup_info))
-        except _store.RestartDmypy:
-            self.api.fail("You probably need to restart dmypy", context)
+        except _store.RestartDmypy as err:
+            self.api.fail(f"You probably need to restart dmypy: {err}", context)
             return AnyType(TypeOfAny.from_error)
         except _store.UnionMustBeOfTypes:
             self.api.fail("Union must be of instances of models", context)

pyproject.toml

@@ -15,7 +15,7 @@ authors = [
 [project.optional-dependencies]
 stubs-latest = [
     "mypy==1.10.0",
-    "django-stubs==5.0.0",
+    "django-stubs==5.0.2",
 ]
 stubs-older = [
     "mypy==1.4.0",

scripts/test_helpers/extended_mypy_django_plugin_test_driver/output_builder.py

@@ -47,6 +47,14 @@ def __init__(
 
         self.target_file = target_file
 
+    def _normalise_message(self, message: str) -> str:
+        if importlib.metadata.version("mypy") == "1.4.0":
+            return message.replace("type[", "Type[").replace(
+                "django.db.models.query.QuerySet", "django.db.models.query._QuerySet"
+            )
+        else:
+            return message
+
     def clear(self) -> Self:
         self._build.result.clear()
         return self
@@ -63,6 +71,11 @@ def on(self, path: str) -> Self:
         return self.__class__(build=self._build, target_file=path)
 
     def from_out(self, out: str, regex: bool = False) -> Self:
+        if importlib.metadata.version("mypy") == "1.4.0":
+            out = out.replace("type[", "Type[").replace(
+                "django.db.models.query.QuerySet", "django.db.models.query._QuerySet"
+            )
+
         self._build.result.extend(
             extract_output_matchers_from_out(
                 out, {}, regex=regex, for_daemon=self._build.for_daemon
@@ -71,17 +84,16 @@ def from_out(self, out: str, regex: bool = False) -> Self:
         return self
 
     def add_revealed_type(self, lnum: int, revealed_type: str) -> Self:
-        if importlib.metadata.version("mypy") == "1.4.0":
-            revealed_type = revealed_type.replace("type[", "Type[")
+        revealed_type = self._normalise_message(revealed_type)
+
         assert self.target_file is not None
         self._build.add(
             self.target_file, lnum, None, "note", f'Revealed type is "{revealed_type}"'
         )
         return self
 
     def change_revealed_type(self, lnum: int, message: str) -> Self:
-        if importlib.metadata.version("mypy") == "1.4.0":
-            message = message.replace("type[", "Type[")
+        message = self._normalise_message(message)
 
         assert self.target_file is not None
 
@@ -122,11 +134,15 @@ def remove_errors(self, lnum: int) -> Self:
         return self
 
     def add_error(self, lnum: int, error_type: str, message: str) -> Self:
+        message = self._normalise_message(message)
+
         assert self.target_file is not None
         self._build.add(self.target_file, lnum, None, "error", f"{message}  [{error_type}]")
         return self
 
     def remove_from_revealed_type(self, lnum: int, remove: str) -> Self:
+        remove = self._normalise_message(remove)
+
         assert self.target_file is not None
 
         found: list[FileOutputMatcher] = []

tests/test_concrete_annotations.py

@@ -43,7 +43,7 @@ def check_instance_with_type_guard(cls: Parent) -> TypeGuard[Concrete[Parent]]:
                 )
                 .add_revealed_type(
                     10,
-                    "Union[django.db.models.query._QuerySet[myapp.models.Child1, myapp.models.Child1], myapp.models.Child2QuerySet, django.db.models.query._QuerySet[myapp.models.Child3, myapp.models.Child3], django.db.models.query._QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]]",
+                    "Union[django.db.models.query.QuerySet[myapp.models.Child1, myapp.models.Child1], myapp.models.Child2QuerySet, django.db.models.query.QuerySet[myapp.models.Child3, myapp.models.Child3], django.db.models.query.QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]]",
                 )
                 .add_revealed_type(
                     20,
@@ -83,7 +83,7 @@ def _(expected: OutputBuilder) -> None:
                 )
                 .add_revealed_type(
                     9,
-                    "Union[django.db.models.query._QuerySet[myapp.models.Child1, myapp.models.Child1], myapp.models.Child2QuerySet, django.db.models.query._QuerySet[myapp.models.Child3, myapp.models.Child3], django.db.models.query._QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]]",
+                    "Union[django.db.models.query.QuerySet[myapp.models.Child1, myapp.models.Child1], myapp.models.Child2QuerySet, django.db.models.query.QuerySet[myapp.models.Child3, myapp.models.Child3], django.db.models.query.QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]]",
                 )
             )
 
@@ -101,7 +101,7 @@ def _(expected: OutputBuilder) -> None:
                 )
                 .remove_from_revealed_type(
                     9,
-                    ", django.db.models.query._QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]",
+                    ", django.db.models.query.QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]",
                 )
             )
 

tests/test_errors.py

@@ -84,6 +84,6 @@ def method_with_instance_typeguard(
                     line
                     for line in out.split("\n")
                     if "Only concrete class can be given" not in line
-                ).replace('Revealed type is "type', 'Revealed type is "Type')
+                )
 
             expected.from_out(out)

tests/test_works.yml → tests/test_works.py

@@ -1,19 +1,22 @@
-- case: mypy_path_from_env
-  start: .
-  disable_cache: true
-  out: |
+from extended_mypy_django_plugin_test_driver import OutputBuilder, Scenario
+
+
+def test_works(scenario: Scenario) -> None:
+    out = """
      main:33: note: Revealed type is "Union[django.db.models.manager.Manager[myapp.models.Child1], myapp.models.ManagerFromChild2QuerySet[myapp.models.Child2], django.db.models.manager.Manager[myapp.models.Child3], django.db.models.manager.Manager[myapp2.models.ChildOther]]"
      main:38: note: Revealed type is "myapp.models.Child1"
-     main:41: note: Revealed type is "Union[django.db.models.query._QuerySet[myapp.models.Child1, myapp.models.Child1], myapp.models.Child2QuerySet, django.db.models.query._QuerySet[myapp.models.Child3, myapp.models.Child3], django.db.models.query._QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]]"
-     main:44: note: Revealed type is "django.db.models.query._QuerySet[myapp.models.Child1, myapp.models.Child1]"
+     main:41: note: Revealed type is "Union[django.db.models.query.QuerySet[myapp.models.Child1, myapp.models.Child1], myapp.models.Child2QuerySet, django.db.models.query.QuerySet[myapp.models.Child3, myapp.models.Child3], django.db.models.query.QuerySet[myapp2.models.ChildOther, myapp2.models.ChildOther]]"
+     main:44: note: Revealed type is "django.db.models.query.QuerySet[myapp.models.Child1, myapp.models.Child1]"
      main:47: note: Revealed type is "myapp.models.Child2QuerySet"
      main:48: note: Revealed type is "myapp.models.Child2QuerySet"
      main:49: note: Revealed type is "myapp.models.ManagerFromChild2QuerySet[myapp.models.Child2]"
      main:50: note: Revealed type is "myapp.models.Child2QuerySet[myapp.models.Child2]"
-     main:53: note: Revealed type is "django.db.models.query._QuerySet[myapp.models.Child1, myapp.models.Child1]"
+     main:53: note: Revealed type is "django.db.models.query.QuerySet[myapp.models.Child1, myapp.models.Child1]"
      main:56: note: Revealed type is "myapp.models.Child2QuerySet"
-     main:59: note: Revealed type is "Union[myapp.models.Child2QuerySet, django.db.models.query._QuerySet[myapp.models.Child1, myapp.models.Child1]]"
-  main: |
+     main:59: note: Revealed type is "Union[myapp.models.Child2QuerySet, django.db.models.query.QuerySet[myapp.models.Child1, myapp.models.Child1]]"
+     """
+
+    main = """
     from extended_mypy_django_plugin import Concrete, ConcreteQuerySet, DefaultQuerySet
 
     from myapp.models import Parent, Child1, Child2
@@ -73,3 +76,10 @@ def ones(model: type[Concrete[Parent]]) -> list[str]:
 
     tvqsmult = make_multiple_queryset(Child1)
     reveal_type(tvqsmult)
+    """
+
+    @scenario.run_and_check_mypy_after
+    def _(expected: OutputBuilder) -> None:
+        scenario.make_file("main.py", main)
+
+        expected.from_out(out)