Add hide_queryset_annotations helper

It's useful to be able to tell mypy to ignore the queryset annotations
when we've only added an annotation for a filter applied in that place.

For example when we want to return a non annotated queryset from a
function, but require an annotation within that function to determine
what's in the queryset

Currently when we do this we need to specify those annotations in the
return type of such a function even when they aren't used outside that
function

This is a response to typeddjango/django-stubs#3382

Author: delfick

docs/api/changelog.rst

@@ -10,6 +10,7 @@ Changelog
     * Dropped support for mypy<2.1.0
     * Dropped support for python<3.13
     * Dropped support for django-stubs<6.0.5,django<6.0.5
+    * Added ``extended_mypy_django_plugin.hide_queryset_annotations`` helper
 
 .. _release-0.8.2:
 

docs/api/usage.rst

@@ -193,3 +193,34 @@ This also works on the concrete models themselves:
 
     qs1: models.QuerySet[Concrete1]
     qs2: Concrete2QuerySet
+
+Hiding queryset annotations
+---------------------------
+
+Sometimes it can be useful to create a queryset using an annotation that isn't
+useful to callers.
+
+For example:
+
+.. code-block:: python
+
+    def some_filter(qs: MyQuerySet[MyModel]) -> MyQuerySet[MyModel]:
+        # Imagine a useful example
+        # The problem is that this fails because mypy believes the annotated
+        # type this resolves to should result in a `return-value` error here
+        return qs.annotate(max_value=Max("value")).filter(max_value__gt=10)
+
+In this case we can remove the queryset annotations:
+
+.. code-block:: python
+
+    from extended_django_mypy_plugin import hide_queryset_annotations
+
+    def some_filter(qs: MyQuerySet[MyModel]) -> MyQuerySet[MyModel]:
+        return hide_queryset_annotations(
+            qs.annotate(max_value=Max("value")).filter(max_value__gt=10)
+        )
+
+This helper does nothing at runtime and at static times errors will be produced
+if it's given something that isn't an instance of a queryset with a model we
+can determine.

extended_mypy_django_plugin/__init__.py

@@ -1,3 +1,3 @@
-from .annotations import Concrete, DefaultQuerySet
+from .annotations import Concrete, DefaultQuerySet, hide_queryset_annotations
 
-__all__ = ["Concrete", "DefaultQuerySet"]
+__all__ = ["Concrete", "DefaultQuerySet", "hide_queryset_annotations"]

extended_mypy_django_plugin/_plugin/plugin.py

@@ -1,9 +1,16 @@
 import functools
 from typing import Generic, TypeVar
 
+from mypy import checker
 from mypy.nodes import Import, ImportAll, ImportFrom, MypyFile
 from mypy.options import Options
-from mypy.plugin import AnalyzeTypeContext, MethodContext, ReportConfigContext
+from mypy.plugin import (
+    AnalyzeTypeContext,
+    FunctionSigContext,
+    MethodContext,
+    ReportConfigContext,
+)
+from mypy.types import FunctionLike
 from mypy.types import Type as MypyType
 from mypy_django_plugin import main
 
@@ -48,6 +55,9 @@ class ExtendedMypyStubs(Generic[T_Report], main.NewSemanalDjangoPlugin):
 
     .. autoattribute:: get_method_hook
         :no-index:
+
+    .. autoattribute:: get_function_signature_hook
+        :no-index:
     """
 
     @classmethod
@@ -201,3 +211,27 @@ def choose(
 
         def run(self, ctx: MethodContext) -> MypyType:
             return self.plugin.type_checker.modify_cast_as_concrete(ctx)
+
+    @hook.hook
+    class get_function_signature_hook(
+        PlainHook[protocols.Report, FunctionSigContext, FunctionLike]
+    ):
+        """
+        Used to implement the ``hide_queryset_annotations`` helper.
+
+        We use this hook to change the signature of that function where it is used such that it takes in the type it is
+        given and returns the same type minus any django-stubs queryset annotation information.
+        """
+
+        def choose(
+            self, *, fullname: str, super_hook: hook.MypyHook[FunctionSigContext, FunctionLike]
+        ) -> bool:
+            return fullname == "extended_mypy_django_plugin.annotations.hide_queryset_annotations"
+
+        def run(self, ctx: FunctionSigContext) -> FunctionLike:
+            return self.plugin.type_checker.modify_hide_queryset_annotations(
+                ctx=ctx,
+                django_context=self.plugin.django_context,
+                lookup_fully_qualified=self.plugin.lookup_fully_qualified,
+                scope=None if not isinstance(ctx.api, checker.TypeChecker) else ctx.api.scope,
+            )

extended_mypy_django_plugin/_plugin/type_checker.py

@@ -1,10 +1,18 @@
+import mypy_django_plugin.django.context
+import mypy_django_plugin.lib.helpers
+from mypy import checker_shared
+from mypy.nodes import TypeInfo
 from mypy.plugin import (
     FunctionContext,
+    FunctionSigContext,
     MethodContext,
 )
 from mypy.types import (
     AnyType,
+    FunctionLike,
     Instance,
+    TypeAliasType,
+    TypedDictType,
     TypeOfAny,
     TypeType,
     TypeVarType,
@@ -20,6 +28,75 @@ class TypeChecking:
     def __init__(self, *, make_resolver: protocols.ResolverMaker) -> None:
         self.make_resolver = make_resolver
 
+    def modify_hide_queryset_annotations(
+        self,
+        *,
+        ctx: FunctionSigContext,
+        django_context: mypy_django_plugin.django.context.DjangoContext,
+        lookup_fully_qualified: protocols.LookupFullyQualified,
+        scope: checker_shared.CheckerScope | None = None,
+    ) -> FunctionLike:
+        if len(ctx.args) != 1:
+            ctx.api.fail("hide_queryset_annotations takes only one argument", ctx.context)
+            return ctx.default_signature
+
+        if not ctx.args[0]:
+            ctx.api.fail("Mypy failed to tell us the type of the first argument", ctx.context)
+            return ctx.default_signature
+
+        first_arg = get_proper_type(ctx.api.get_expression_type(ctx.args[0][0]))
+        if isinstance(first_arg, AnyType):
+            ctx.api.fail("Failed to determine the type of the first argument", ctx.context)
+            return ctx.default_signature
+
+        if isinstance(first_arg, TypeAliasType) and first_arg.alias:
+            first_arg = get_proper_type(first_arg.alias.target)
+
+        if isinstance(first_arg, TypeVarType) and first_arg.name == "Self":
+            if (
+                scope is None
+                or (enclosing := scope.enclosing_class()) is None
+                or not enclosing.has_base("django.db.models.query.QuerySet")
+            ):
+                ctx.api.fail("First argument must be a queryset", ctx.context)
+
+            return ctx.default_signature
+
+        if not isinstance(first_arg, Instance):
+            ctx.api.fail("First argument must be an instance", ctx.context)
+            return ctx.default_signature
+
+        model = mypy_django_plugin.lib.helpers.extract_model_type_from_queryset(first_arg)
+        if model is None:
+            ctx.api.fail(
+                "Failed to determine a django model from the first argument (or it's not a queryset)",
+                ctx.context,
+            )
+            return ctx.default_signature
+
+        if mypy_django_plugin.lib.helpers.is_annotated_model(model.type):
+            fullname = model.type.bases[0].type.fullname
+        else:
+            fullname = model.type.fullname
+
+        if (sym := lookup_fully_qualified(fullname)) is None or not isinstance(sym.node, TypeInfo):
+            ctx.api.fail(f"Failed to determine a django model from '{fullname}'", ctx.context)
+            return ctx.default_signature
+
+        django_model = Instance(sym.node, [])
+        queryset_args: list[MypyType] = [django_model]
+        if len(first_arg.args) == 2:
+            second_arg = get_proper_type(first_arg.args[1])
+            if isinstance(second_arg, TypeAliasType) and second_arg.alias:
+                second_arg = get_proper_type(second_arg.alias.target)
+
+            if isinstance(second_arg, TypedDictType):
+                queryset_args.append(first_arg.args[1])
+
+        return ctx.default_signature.copy_modified(
+            arg_types=[first_arg], ret_type=Instance(first_arg.type, queryset_args)
+        )
+
     def modify_cast_as_concrete(self, ctx: FunctionContext | MethodContext) -> MypyType:
         if len(ctx.arg_types) != 1:
             ctx.api.fail("Concrete.cast_as_concrete takes only one argument", ctx.context)

extended_mypy_django_plugin/annotations.py

@@ -8,6 +8,33 @@
 T_Obj = TypeVar("T_Obj")
 
 
+def hide_queryset_annotations[T_QuerySet](queryset: T_QuerySet) -> T_QuerySet:
+    """
+    This helper is seen by the mypy plugin and is used to transform a queryset
+    type with annotations into a queryset type without annotations
+
+    At runtime it does nothing.
+
+    The purpose of this is when a queryset is annotated only for the filter that is
+    immediately applied to it when we don't need to share the annotation outside of
+    that context.
+
+    For example:
+
+    .. code-block::
+
+        from extended_mypy_django_plugin import hide_queryset_annotations
+
+        def filter_my_queryset(qs: MySpecialQuerySet[MyModel]) -> MySpecialQuerySet[MyModel]:
+            # Without the `hide_queryset_annotations` mypy will complain that we are returning
+            # MySpecialQuerySet[Annotated[MyModel, TypedDict({"some_field": ...})]]
+            return hide_queryset_annotations(
+                qs.annotate(some_field=...).filter(some_field__gt=3)
+            )
+    """
+    return queryset
+
+
 class Concrete(Generic[T_Parent]):
     """
     The ``Concrete`` annotation exists as a class with functionality for both

scripts/follower1/models/__init__.py

@@ -1,3 +1,3 @@
-from .follower1 import Follower1
+from .follower1 import Follower1, Follower1QuerySet
 
-__all__ = ["Follower1"]
+__all__ = ["Follower1", "Follower1QuerySet"]