From 2720b74242380e0138d8e5e75cb2e6f4310e606b Mon Sep 17 00:00:00 2001 From: Maxim Kurnikov Date: Wed, 13 Feb 2019 14:39:52 +0300 Subject: [PATCH] add proper generic support for get_object_or_404/get_list_or_404, fixes #22 --- django-stubs/shortcuts.pyi | 11 +++++++---- scripts/typecheck_tests.py | 7 +++---- test-data/typecheck/managers.test | 15 +++++++++++++++ 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/django-stubs/shortcuts.pyi b/django-stubs/shortcuts.pyi index f245745..7cd19a9 100644 --- a/django-stubs/shortcuts.pyi +++ b/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: ... diff --git a/scripts/typecheck_tests.py b/scripts/typecheck_tests.py index 197157d..3e94e94 100644 --- a/scripts/typecheck_tests.py +++ b/scripts/typecheck_tests.py @@ -149,11 +149,10 @@ IGNORED_ERRORS = { ], '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[], Manager[], QuerySet[]]"', 'Argument 1 to "get_list_or_404" has incompatible type "List[Type[Article]]"; ' - + 'expected "Union[Type[Model], Manager[Any], QuerySet[Any]]"' + + 'expected "Union[Type[], Manager[], QuerySet[]]"', + 'CustomClass' ], 'get_or_create': [ 'Argument 1 to "update_or_create" of "QuerySet" has incompatible type "**Dict[str, object]"; expected "MutableMapping[str, Any]"' diff --git a/test-data/typecheck/managers.test b/test-data/typecheck/managers.test index 88ad066..aa96538 100644 --- a/test-data/typecheck/managers.test +++ b/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] \ No newline at end of file