update add_ordering signature (#2822)

Author: UnknownPlatypus

django-stubs/db/models/query.pyi

@@ -24,6 +24,8 @@ _PrefetchedQuerySetT = TypeVar("_PrefetchedQuerySetT", bound=QuerySet[Model], co
 # This will be specialized to a `LiteralString` in the plugin for further processing and validation
 _ToAttrT = TypeVar("_ToAttrT", bound=str, covariant=True, default=str)
 
+_OrderByFieldName: TypeAlias = str | Combinable
+
 MAX_GET_RESULTS: int
 REPR_OUTPUT_SIZE: int
 
@@ -189,7 +191,7 @@ class QuerySet(Generic[_Model, _Row], Iterable[_Row], Sized):
     def prefetch_related(self, *lookups: str | Prefetch[_LookupT, _PrefetchedQuerySetT, _ToAttrT]) -> Self: ...
     def annotate(self, *args: Any, **kwargs: Any) -> Self: ...
     def alias(self, *args: Any, **kwargs: Any) -> Self: ...
-    def order_by(self, *field_names: str | Combinable) -> Self: ...
+    def order_by(self, *field_names: _OrderByFieldName) -> Self: ...
     def distinct(self, *field_names: str) -> Self: ...
     # extra() return type won't be supported any time soon
     def extra(
@@ -198,7 +200,7 @@ class QuerySet(Generic[_Model, _Row], Iterable[_Row], Sized):
         where: Sequence[str] | None = None,
         params: Sequence[Any] | None = None,
         tables: Sequence[str] | None = None,
-        order_by: Sequence[str] | None = None,
+        order_by: Sequence[_OrderByFieldName] | None = None,
         select_params: Sequence[Any] | None = None,
     ) -> QuerySet[Any, Any]: ...
     def reverse(self) -> Self: ...

django-stubs/db/models/sql/query.pyi

@@ -4,9 +4,10 @@ from typing import Any, Literal, NamedTuple
 
 from django.db.backends.utils import CursorWrapper
 from django.db.models import Field, FilteredRelation, Model, Q
-from django.db.models.expressions import BaseExpression, Combinable, Expression, OrderBy
+from django.db.models.expressions import BaseExpression, Combinable, Expression
 from django.db.models.lookups import Lookup, Transform
 from django.db.models.options import Options
+from django.db.models.query import _OrderByFieldName
 from django.db.models.query_utils import PathInfo
 from django.db.models.sql.datastructures import BaseTable, Join
 from django.db.models.sql.where import WhereNode
@@ -58,7 +59,7 @@ class Query(BaseExpression):
     filter_is_sticky: bool
     subquery: bool
     group_by: None | Sequence[Combinable] | Sequence[str] | Literal[True]
-    order_by: Sequence[Any]
+    order_by: Sequence[_OrderByFieldName]
     distinct: bool
     distinct_fields: tuple[str, ...]
     select: Sequence[BaseExpression]
@@ -78,7 +79,7 @@ class Query(BaseExpression):
     combined_queries: tuple
     extra_select_mask: set[str] | None
     extra_tables: tuple
-    extra_order_by: Sequence[Any]
+    extra_order_by: Sequence[_OrderByFieldName]
     deferred_loading: tuple[set[str] | frozenset[str], bool]
     explain_query: bool
     explain_format: str | None
@@ -190,7 +191,7 @@ class Query(BaseExpression):
     def set_select(self, cols: list[Expression]) -> None: ...
     def add_distinct_fields(self, *field_names: str) -> None: ...
     def add_fields(self, field_names: Iterable[str], allow_m2m: bool = True) -> None: ...
-    def add_ordering(self, *ordering: str | OrderBy) -> None: ...
+    def add_ordering(self, *ordering: _OrderByFieldName) -> None: ...
     def clear_where(self) -> None: ...
     def clear_ordering(self, force: bool = False, clear_default: bool = True) -> None: ...
     def set_group_by(self, allow_aliases: bool = True) -> None: ...
@@ -202,7 +203,7 @@ class Query(BaseExpression):
         where: Sequence[str] | None,
         params: Sequence[str] | None,
         tables: Sequence[str] | None,
-        order_by: Sequence[str] | None,
+        order_by: Sequence[_OrderByFieldName] | None,
     ) -> None: ...
     def clear_deferred_loading(self) -> None: ...
     def add_deferred_loading(self, field_names: Iterable[str]) -> None: ...

tests/assert_type/db/models/test_ordering.py

@@ -0,0 +1,55 @@
+from django.contrib.auth.models import User
+from django.db.models import F, OrderBy
+from django.db.models.sql.query import Query
+
+qs = User.objects.all()
+qs.order_by("username")
+qs.order_by(F("username"))
+qs.order_by("-date_joined", F("last_login").desc())
+
+# Test OrderBy expressions (success cases)
+qs.order_by(OrderBy(F("username")))
+qs.order_by(OrderBy(F("date_joined"), descending=True))
+qs.order_by("username", OrderBy(F("last_login"), descending=True))
+
+qs.extra(order_by=["username"])
+qs.extra(order_by=[F("username")])
+qs.extra(order_by=["username", F("date_joined").desc()])
+qs.extra(order_by=[OrderBy(F("username"))])
+qs.extra(order_by=["username", OrderBy(F("date_joined"), descending=True)])
+
+query = Query(User)
+query.add_ordering("username")
+query.add_ordering(F("date_joined"))
+query.add_ordering("-last_login", F("username").desc())
+query.add_ordering(OrderBy(F("username")))
+query.add_ordering("username", OrderBy(F("date_joined"), descending=True))
+
+query.add_extra(
+    select=None,
+    select_params=None,
+    where=None,
+    params=None,
+    tables=None,
+    order_by=["username", F("date_joined").desc(), OrderBy(F("last_login"))],
+)
+
+# failure cases
+qs.order_by(123)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+qs.order_by(["username"])  # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+qs.order_by({"username": "asc"})  # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+
+qs.extra(order_by=[123])  # type: ignore[list-item] # pyright: ignore[reportArgumentType]
+qs.extra(order_by=["username", 456])  # type: ignore[list-item] # pyright: ignore[reportArgumentType]
+
+query.add_ordering(123)  # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+query.add_ordering(["username"])  # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+
+query.add_extra(
+    select=None,
+    select_params=None,
+    where=None,
+    params=None,
+    tables=None,
+    order_by=[123, "username"],  # type: ignore[list-item] # pyright: ignore[reportArgumentType]
+)

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -668,7 +668,7 @@
         reveal_type(MyModel.objects.difference)  # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]) -> myapp.models.MyQuerySet"
         reveal_type(MyModel.objects.distinct)  # N: Revealed type is "def (*field_names: builtins.str) -> myapp.models.MyQuerySet"
         reveal_type(MyModel.objects.exclude)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
-        reveal_type(MyModel.objects.extra)  # N: Revealed type is "def (select: builtins.dict[builtins.str, Any] | None =, where: typing.Sequence[builtins.str] | None =, params: typing.Sequence[Any] | None =, tables: typing.Sequence[builtins.str] | None =, order_by: typing.Sequence[builtins.str] | None =, select_params: typing.Sequence[Any] | None =) -> myapp.models.MyQuerySet"
+        reveal_type(MyModel.objects.extra)  # N: Revealed type is "def (select: builtins.dict[builtins.str, Any] | None =, where: typing.Sequence[builtins.str] | None =, params: typing.Sequence[Any] | None =, tables: typing.Sequence[builtins.str] | None =, order_by: typing.Sequence[builtins.str | django.db.models.expressions.Combinable] | None =, select_params: typing.Sequence[Any] | None =) -> myapp.models.MyQuerySet"
         reveal_type(MyModel.objects.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
         reveal_type(MyModel.objects.intersection)  # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]) -> myapp.models.MyQuerySet"
         reveal_type(MyModel.objects.none)  # N: Revealed type is "def () -> myapp.models.MyQuerySet"