add proper generic support for get_object_or_404/get_list_or_404, fixes #22

Author: mkurnikov

django-stubs/shortcuts.pyi

@@ -1,9 +1,9 @@
-from typing import Any, Callable, Dict, List, Optional, Type, Union, Sequence, Protocol
+from typing import Any, Callable, Dict, List, Optional, Protocol, Sequence, Type, TypeVar, Union
 
-from django.db.models import Manager, QuerySet
 from django.db.models.base import Model
 from django.http.response import HttpResponse as HttpResponse, HttpResponseRedirect as HttpResponseRedirect
 
+from django.db.models import Manager, QuerySet
 from django.http import HttpRequest
 
 def render_to_response(
@@ -28,6 +28,9 @@ class SupportsGetAbsoluteUrl(Protocol):
 def redirect(
     to: Union[Callable, str, SupportsGetAbsoluteUrl], *args: Any, permanent: bool = ..., **kwargs: Any
 ) -> HttpResponseRedirect: ...
-def get_object_or_404(klass: Union[Type[Model], Manager, QuerySet], *args: Any, **kwargs: Any) -> Model: ...
-def get_list_or_404(klass: Union[Type[Model], Manager, QuerySet], *args: Any, **kwargs: Any) -> List[Model]: ...
+
+_T = TypeVar("_T", bound=Model)
+
+def get_object_or_404(klass: Union[Type[_T], Manager[_T], QuerySet[_T]], *args: Any, **kwargs: Any) -> _T: ...
+def get_list_or_404(klass: Union[Type[_T], Manager[_T], QuerySet[_T]], *args: Any, **kwargs: Any) -> List[_T]: ...
 def resolve_url(to: Union[Callable, Model, str], *args: Any, **kwargs: Any) -> str: ...

scripts/typecheck_tests.py

@@ -149,11 +149,10 @@
     ],
     'get_object_or_404': [
         'Argument 1 to "get_object_or_404" has incompatible type "str"; '
-        + 'expected "Union[Type[Model], Manager[Any], QuerySet[Any]]"',
-        'Argument 1 to "get_object_or_404" has incompatible type "Type[CustomClass]"; '
-        + 'expected "Union[Type[Model], Manager[Any], QuerySet[Any]]"',
+        + 'expected "Union[Type[<nothing>], Manager[<nothing>], QuerySet[<nothing>]]"',
         'Argument 1 to "get_list_or_404" has incompatible type "List[Type[Article]]"; '
-        + 'expected "Union[Type[Model], Manager[Any], QuerySet[Any]]"'
+        + 'expected "Union[Type[<nothing>], Manager[<nothing>], QuerySet[<nothing>]]"',
+        'CustomClass'
     ],
     'get_or_create': [
         'Argument 1 to "update_or_create" of "QuerySet" has incompatible type "**Dict[str, object]"; expected "MutableMapping[str, Any]"'

test-data/typecheck/managers.test

@@ -136,4 +136,19 @@ class AbstractBase2(models.Model):
 
 class Child(AbstractBase1, AbstractBase2):
     pass
+[out]
+
+[CASE get_object_or_404_returns_proper_types]
+from django.shortcuts import get_object_or_404, get_list_or_404
+from django.db import models
+
+class MyModel(models.Model):
+    pass
+reveal_type(get_object_or_404(MyModel))  # E: Revealed type is 'main.MyModel*'
+reveal_type(get_object_or_404(MyModel.objects))  # E: Revealed type is 'main.MyModel*'
+reveal_type(get_object_or_404(MyModel.objects.get_queryset()))  # E: Revealed type is 'main.MyModel*'
+
+reveal_type(get_list_or_404(MyModel))  # E: Revealed type is 'builtins.list[main.MyModel*]'
+reveal_type(get_list_or_404(MyModel.objects))  # E: Revealed type is 'builtins.list[main.MyModel*]'
+reveal_type(get_list_or_404(MyModel.objects.get_queryset()))  # E: Revealed type is 'builtins.list[main.MyModel*]'
 [out]