Improve narrowing return types

docs/api/changelog.rst

@@ -3,6 +3,13 @@
 Changelog
 ---------
 
+.. _release-0.5.4:
+
+0.5.4 - TBD
+    * Will now check return types for methods and functions more thorouhgly
+    * Will throw errors if a type guard is used with a concrete annotation that uses
+      a type var (mypy plugin system is limited in a way that makes this impossible to implement)
+
 .. _release-0.5.3:
 
 0.5.3 - 25 May 2024

extended_mypy_django_plugin/plugin/_known_annotations.py

@@ -0,0 +1,7 @@
+import enum
+
+
+class KnownAnnotations(enum.Enum):
+    CONCRETE = "extended_mypy_django_plugin.annotations.Concrete"
+    CONCRETE_QUERYSET = "extended_mypy_django_plugin.annotations.ConcreteQuerySet"
+    DEFAULT_QUERYSET = "extended_mypy_django_plugin.annotations.DefaultQuerySet"

extended_mypy_django_plugin/plugin/_plugin.py

@@ -1,5 +1,4 @@
 import collections
-import enum
 import sys
 from typing import Generic
 
@@ -13,10 +12,13 @@
     ClassDefContext,
     DynamicClassDefContext,
     FunctionContext,
+    FunctionSigContext,
+    MethodContext,
+    MethodSigContext,
 )
 from mypy.semanal import SemanticAnalyzer
 from mypy.typeanal import TypeAnalyser
-from mypy.types import CallableType, Instance
+from mypy.types import CallableType, FunctionLike, Instance
 from mypy.types import Type as MypyType
 from mypy_django_plugin import main
 from mypy_django_plugin.django.context import DjangoContext
@@ -26,7 +28,16 @@
 )
 from typing_extensions import assert_never
 
-from . import _config, _dependencies, _helpers, _hook, _reports, _store, actions
+from . import (
+    _config,
+    _dependencies,
+    _helpers,
+    _hook,
+    _known_annotations,
+    _reports,
+    _store,
+    actions,
+)
 
 
 class Hook(
@@ -62,10 +73,7 @@ class ExtendedMypyStubs(main.NewSemanalDjangoPlugin):
 
     plugin_config: _config.Config
 
-    class Annotations(enum.Enum):
-        CONCRETE = "extended_mypy_django_plugin.annotations.Concrete"
-        CONCRETE_QUERYSET = "extended_mypy_django_plugin.annotations.ConcreteQuerySet"
-        DEFAULT_QUERYSET = "extended_mypy_django_plugin.annotations.DefaultQuerySet"
+    Annotations = _known_annotations.KnownAnnotations
 
     def __init__(self, options: Options, mypy_version_tuple: tuple[int, int]) -> None:
         super(main.NewSemanalDjangoPlugin, self).__init__(options)
@@ -263,19 +271,57 @@ def run(self, ctx: AnalyzeTypeContext) -> MypyType:
 
             elif name is Known.DEFAULT_QUERYSET:
                 method = type_analyzer.find_default_queryset
+
             else:
                 assert_never(name)
 
             return method(unbound_type=ctx.type)
 
     @_hook.hook
-    class get_function_hook(Hook[FunctionContext, MypyType]):
+    class get_attribute_hook(Hook[AttributeContext, MypyType]):
         """
-        Find functions that return a ``DefaultQuerySet`` annotation with a type variable
-        and resolve the annotation.
+        An implementation of the change found in
+        https://github.com/typeddjango/django-stubs/pull/2027
         """
 
         def choose(self) -> bool:
+            return self.super_hook is resolve_manager_method
+
+        def run(self, ctx: AttributeContext) -> MypyType:
+            assert isinstance(ctx.api, TypeChecker)
+
+            type_checking = actions.TypeChecking(
+                self.store, api=ctx.api, lookup_info=self.plugin._lookup_info
+            )
+
+            return type_checking.extended_get_attribute_resolve_manager_method(
+                ctx, resolve_manager_method_from_instance=resolve_manager_method_from_instance
+            )
+
+    class SharedCallableHookLogic:
+        """
+        Shared logic for modifying the return type of methods and functions that use a concrete
+        annotation with a type variable.
+
+        Note that the signature hook will already raise errors if a concrete annotation is
+        used with a type var in a type guard.
+        """
+
+        def __init__(self, fullname: str, plugin: "ExtendedMypyStubs") -> None:
+            self.plugin = plugin
+            self.store = plugin.store
+            self.fullname = fullname
+
+        def choose(self) -> bool:
+            """
+            Choose methods and functions either returning a type guard or have a generic
+            return type.
+
+            We determine whether the return type is a concrete annotation or not in the run method.
+            """
+            if self.fullname.startswith("builtins."):
+                return False
+
             sym = self.plugin.lookup_fully_qualified(self.fullname)
             if not sym or not sym.node:
                 return False
@@ -284,18 +330,31 @@ def choose(self) -> bool:
             if not isinstance(call, CallableType):
                 return False
 
-            return call.is_generic()
+            return bool(call.type_guard or call.is_generic())
 
-        def run(self, ctx: FunctionContext) -> MypyType:
+        def run(self, ctx: MethodContext | FunctionContext) -> MypyType | None:
             assert isinstance(ctx.api, TypeChecker)
 
-            type_checking = actions.TypeChecking(self.store, api=ctx.api)
+            type_checking = actions.TypeChecking(
+                self.store,
+                api=ctx.api,
+                lookup_info=self.plugin._lookup_info,
+            )
+
+            return type_checking.modify_return_type(ctx)
 
-            result = type_checking.modify_default_queryset_return_type(
-                ctx,
-                desired_annotation_fullname=ExtendedMypyStubs.Annotations.DEFAULT_QUERYSET.value,
+    @_hook.hook
+    class get_method_hook(Hook[MethodContext, MypyType]):
+        def extra_init(self) -> None:
+            self.shared_logic = self.plugin.SharedCallableHookLogic(
+                fullname=self.fullname, plugin=self.plugin
             )
 
+        def choose(self) -> bool:
+            return self.shared_logic.choose()
+
+        def run(self, ctx: MethodContext) -> MypyType:
+            result = self.shared_logic.run(ctx)
             if result is not None:
                 return result
 
@@ -305,20 +364,106 @@ def run(self, ctx: FunctionContext) -> MypyType:
             return ctx.default_return_type
 
     @_hook.hook
-    class get_attribute_hook(Hook[AttributeContext, MypyType]):
+    class get_function_hook(Hook[FunctionContext, MypyType]):
+        def extra_init(self) -> None:
+            self.shared_logic = self.plugin.SharedCallableHookLogic(
+                fullname=self.fullname, plugin=self.plugin
+            )
+
+        def choose(self) -> bool:
+            return self.shared_logic.choose()
+
+        def run(self, ctx: FunctionContext) -> MypyType:
+            result = self.shared_logic.run(ctx)
+
+            if result is not None:
+                return result
+
+            if self.super_hook is not None:
+                return self.super_hook(ctx)
+
+            return ctx.default_return_type
+
+    class SharedSignatureHookLogic:
         """
-        An implementation of the change found in
-        https://github.com/typeddjango/django-stubs/pull/2027
+        Shared logic for modifying the signature of methods and functions.
+
+        This is only used to find cases where a concrete annotation with a type var
+        is used in a type guard.
+
+        In this situation the mypy plugin system does not provide an opportunity to fully resolve
+        the type guard.
         """
 
+        def __init__(self, fullname: str, plugin: "ExtendedMypyStubs") -> None:
+            self.plugin = plugin
+            self.store = plugin.store
+            self.fullname = fullname
+
         def choose(self) -> bool:
-            return self.super_hook is resolve_manager_method
+            """
+            Only choose methods and functions that are returning a type guard
+            """
+            if self.fullname.startswith("builtins."):
+                return False
 
-        def run(self, ctx: AttributeContext) -> MypyType:
+            sym = self.plugin.lookup_fully_qualified(self.fullname)
+            if not sym or not sym.node:
+                return False
+
+            call = getattr(sym.node, "type", None)
+            if not isinstance(call, CallableType):
+                return False
+
+            return call.type_guard is not None
+
+        def run(self, ctx: MethodSigContext | FunctionSigContext) -> MypyType | None:
             assert isinstance(ctx.api, TypeChecker)
 
-            type_checking = actions.TypeChecking(self.store, api=ctx.api)
+            type_checking = actions.TypeChecking(
+                self.store,
+                api=ctx.api,
+                lookup_info=self.plugin._lookup_info,
+            )
 
-            return type_checking.extended_get_attribute_resolve_manager_method(
-                ctx, resolve_manager_method_from_instance=resolve_manager_method_from_instance
+            return type_checking.check_typeguard(ctx.context)
+
+    @_hook.hook
+    class get_function_signature_hook(Hook[FunctionSigContext, FunctionLike]):
+        def extra_init(self) -> None:
+            self.shared_logic = self.plugin.SharedSignatureHookLogic(
+                fullname=self.fullname, plugin=self.plugin
+            )
+
+        def choose(self) -> bool:
+            return self.shared_logic.choose()
+
+        def run(self, ctx: FunctionSigContext) -> FunctionLike:
+            result = self.shared_logic.run(ctx)
+            if result is not None:
+                return ctx.default_signature.copy_modified(ret_type=result)
+
+            if self.super_hook is not None:
+                return self.super_hook(ctx)
+
+            return ctx.default_signature
+
+    @_hook.hook
+    class get_method_signature_hook(Hook[MethodSigContext, FunctionLike]):
+        def extra_init(self) -> None:
+            self.shared_logic = self.plugin.SharedSignatureHookLogic(
+                fullname=self.fullname, plugin=self.plugin
             )
+
+        def choose(self) -> bool:
+            return self.shared_logic.choose()
+
+        def run(self, ctx: MethodSigContext) -> FunctionLike:
+            result = self.shared_logic.run(ctx)
+            if result is not None:
+                return ctx.default_signature.copy_modified(ret_type=result)
+
+            if self.super_hook is not None:
+                return self.super_hook(ctx)
+
+            return ctx.default_signature

extended_mypy_django_plugin/plugin/actions/_type_analyze.py

@@ -30,7 +30,7 @@ def find_concrete_models(self, unbound_type: UnboundType) -> MypyType:
         type_arg = get_proper_type(self.api.analyze_type(args[0]))
 
         if not isinstance(type_arg, Instance):
-            return UnionType(())
+            return unbound_type
 
         concrete = tuple(
             self.store.retrieve_concrete_children_types(