django-stubs icon indicating copy to clipboard operation
django-stubs copied to clipboard

custom queryset methods/types get lost following calls to `.values()`, `.values_list()`

Open bwo opened this issue 2 years ago • 2 comments

(and possibly other methods; in fact, I would bet any method that isn't annotated as returning Self in query.pyi.)

#$ cat foo.py 
from typing import Self, TypeVar
from django.db import models

T = TypeVar("T", bound=models.Model)

class MyQuerySet(models.QuerySet[T]):
    def filter_field(self, x: list[str]) -> Self:
        return self.filter(field__in=x)

mgr = models.Manager.from_queryset(MyQuerySet)


class MyModelA(models.Model):
    # objects = models.Manager()
    objects = mgr()

    field = models.CharField()


reveal_type(MyModelA.objects.filter(id__gt=1).filter_field(["hello"]).first())
reveal_type(MyModelA.objects.filter(id__gt=1).values("field").filter_field(["hello"]).first())  # AttributeError. typed as Any
reveal_type(MyModelA.objects.filter(id__gt=1).filter_field(["hello"]).values("id").first())

Then running mypy:

foo.py:20: note: Revealed type is "Union[payment.models.MyModelA, None]"
foo.py:21 error: "_QuerySet[MyModelA, TypedDict({'field': str})]" has no attribute "filter_field"  [attr-defined]
foo.py:21: note: Revealed type is "Any"
foo.py:22: note: Revealed type is "Union[TypedDict({'id': builtins.int}), None]"

That is: the chaining works following filter, and the chaining works if .values() is called after a custom method. But it doesn't work if .values() is called first—note that it's lost the belief that the queryset is of type MyQuerySet.

I believe (as noted) that the root cause here is that .values() is annotated in query.pyi like this:

    def values(self, *fields: str | Combinable, **expressions: Any) -> _QuerySet[_T, dict[str, Any]]: ... 

So of course it loses track of the actual type! I'm not sure if there's a good (or any) way to notate "the same type constructor as me, but differently parameterized" in mypy, though.

System information

  • OS: macos 13.5.2
  • python version: 3.11
  • django version: 4.2
  • mypy version: 1.6.1 (also happens with 1.7 and 1.8-dev versions)
  • django-stubs version: 4.2.6
  • django-stubs-ext version: 4.2.5

bwo avatar Nov 16 '23 20:11 bwo

note that it's lost the belief that the queryset is of type MyQuerySet.

This is correct. .values() returns a QuerySet[dict]. In your example it returns QuerySet[T] where T = TypeVar("T", bound=models.Model), the upper bound here not being compatible with dict.

Try using ValuesQuerySet instead. It is sortof a wrapper for QuerySet[dict]

hterik avatar Sep 30 '24 12:09 hterik

How would using something different in my code help with the type inference guided by query.pyi? The problem is that values() is annotated as returning _Queryset[...], rather than (asfilter() does) Self. It's kind of between a rock and a hard place because Self is fully specified (you can't return Self[_T, dict[str, Any]]), but this seems like something that, perhaps, the plugin could patch up, just as the plugin synthesizes a TypedDict rather than just using a dict[str, Any].

bwo avatar Sep 30 '24 12:09 bwo