models: Tighten type annotation for query_for_ids.

This makes use of the type parameters of ValuesQuerySet (a.k.a
_QuerySet) to provide a more accurate type annotation for query_for_ids.

Note that QuerySet[ModelT] is equivalent to _QuerySet[ModelT, ModelT].

Signed-off-by: Zixuan James Li <p359101898@gmail.com>
This commit is contained in:
Zixuan James Li 2022-06-24 11:53:28 -04:00 committed by Anders Kaseorg
parent 5441ec7d16
commit 3400cb425b
1 changed files with 9 additions and 1 deletions

View File

@ -155,7 +155,15 @@ class AndNonZero(models.Lookup):
return f"{lhs} & {rhs} != 0", lhs_params + rhs_params
def query_for_ids(query: QuerySet, user_ids: List[int], field: str) -> QuerySet:
ModelT = TypeVar("ModelT", bound=models.Model)
RowT = TypeVar("RowT")
def query_for_ids(
query: "ValuesQuerySet[ModelT, RowT]",
user_ids: List[int],
field: str,
) -> "ValuesQuerySet[ModelT, RowT]":
"""
This function optimizes searches of the form
`user_profile_id in (1, 2, 3, 4)` by quickly