Support returning the correct values for the different QuerySet methods when using .values() and .values_list().

[syastrov]

No description provided.

[mkurnikov]

It's not clear to me, how all of that works. Could you add tests for other methods of QuerySet too? I'd split it into a number of tests:

1. first(), earliest(), __getitem__, etc. for plain QuerySet.
2. Same but for values() (this should return TypedDict some time in the future)
3. Same but for values_list().

in_bulk may or may not be as part of every test, add it where it's applicable.

[syastrov]

I have added a lot more tests, and in the process, caught that I had missed updating the .iterator method.

I haven't tested all of the queryset methods that I changed at the moment, as it's quite tedious to create all the cases. If I get some more time, I will add some more.

The annoying thing is that you cannot tell mypy that the default type for _Row is _T. It would make it a lot more user-friendly for annotating your own code.
Unfortunately, right now I believe you would have to e.g. write QuerySet[MyModel, MyModel] to avoid having mypy type it as QuerySet[MyModel, Any].
They have an open issue for that: python/mypy#4236

[mkurnikov]

I haven't tested all of the queryset methods that I changed at the moment, as it's quite tedious to create all the cases. If I get some more time, I will add some more.

Sorry, I was wrong/imprecise in my comment, testing all the methods is indeed very tedious and pretty redundant. It should be enough to select just one of the group with the same number of generics present in the signature and same return type, like one of

    def none(self) -> QuerySet[_T]: ...
    def all(self) -> QuerySet[_T]: ...
    def filter(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
    def exclude(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
    def complex_filter(self, filter_obj: Any) -> QuerySet[_T]: ...
    def count(self) -> int: ...
    def union(self, *other_qs: Any, all: bool = ...) -> QuerySet[_T]: ...
    def intersection(self, *other_qs: Any) -> QuerySet[_T]: ...
    def difference(self, *other_qs: Any) -> QuerySet[_T]: ...
    def select_for_update(self, nowait: bool = ..., skip_locked: bool = ..., of: Tuple = ...) -> QuerySet: ...
    def select_related(self, *fields: Any) -> QuerySet[_T]: ...
    def prefetch_related(self, *lookups: Any) -> QuerySet[_T]: ...
    def annotate(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
    def order_by(self, *field_names: Any) -> QuerySet[_T]: ...
    def distinct(self, *field_names: Any) -> QuerySet[_T]: ...
    def extra(
        self,
        select: Optional[Dict[str, Any]] = ...,
        where: Optional[List[str]] = ...,
        params: Optional[List[Any]] = ...,
        tables: Optional[List[str]] = ...,
        order_by: Optional[Sequence[str]] = ...,
        select_params: Optional[Sequence[Any]] = ...,
    ) -> QuerySet[_T]: ...
    def reverse(self) -> QuerySet[_T]: ...
    def defer(self, *fields: Any) -> QuerySet[_T]: ...
    def only(self, *fields: Any) -> QuerySet[_T]: ...
    def using(self, alias: Optional[str]) -> QuerySet[_T]: ...

one of

    def earliest(self, *fields: Any, field_name: Optional[Any] = ...) -> _T: ...
    def latest(self, *fields: Any, field_name: Optional[Any] = ...) -> _T: ...
    def first(self) -> Optional[_T]: ...
    def last(self) -> Optional[_T]: ...

and so on.

And thank you very much for your efforts!

[mkurnikov]

Unfortunately, right now I believe you would have to e.g. write QuerySet[MyModel, MyModel] to avoid having mypy type it as QuerySet[MyModel, Any].

We can hide it from users by leveragive get_type_analyze_hook:

def set_first_generic_param_as_default_for_second(ctx: AnalyzeTypeContext) -> Type:
    if not ctx.type.args:
        return ctx.api.named_type(helpers.QUERYSET_CLASS_FULLNAME, [AnyType(TypeOfAny.explicit),
                                                                    AnyType(TypeOfAny.explicit)])
    args = ctx.type.args
    if len(args) == 1:
        args = [args[0], args[0]]

    analyzed_args = [ctx.api.analyze_type(arg) for arg in args]
    # TODO: this function should be made into an object with __call__, and initialized with fullname,
    # s/helpers.QUERYSET_CLASS_FULLNAME/fullname
    return ctx.api.named_type(helpers.QUERYSET_CLASS_FULLNAME, analyzed_args)

def get_type_analyze_hook(self, fullname: str
                              ) -> Optional[Callable[[AnalyzeTypeContext], Type]]:
        # TODO: instead of just queryset, it should be in the same style as _get_current_model_bases/_get_current_manager_bases/...
        if fullname == 'django.db.models.query.QuerySet':
            return set_first_generic_param_as_default_for_second

Please, address TODOs, see also python/mypy#6503.

Error messages should also be changed, but let's do it some time later, as it should also be done for all Field subclasses, see python/mypy#6501

[syastrov]

That sounds like a good plan by using get_type_analyze_hook. I'll see if I get the time to implement this soon.

[syastrov]

I think it's ready for your review again now :)

I made a few more changes to fix some small issues I noticed as well:

• Fix return type for QuerySet.select_for_update() (missing type params)
• I have added correct return type for QuerySet.dates() / QuerySet.datetimes()
• Added correct type params in return type for QuerySet.__and__ / QuerySet.__or__

[mkurnikov]

PR looks great, resolve conflicts with master again and I'll merge.