Improve stubs for flatten and flatten_fieldsets (#2572)

* Add new `_ListDisplayT` type alias

This renames `_DisplayT` to `_ListDisplayT` and creates a new
`_DisplayT` alias with the contents of the `_ListOrTuple`.

This will be used in the next commit.

* Improve stubs for `flatten` and `flatten_fieldsets`

These functions are implemented in a generic way that doesn't
generally care about the content. The existing return type of
`Callable | str` makes these awkward to use.

I've changed the use of `Sequence` in the `_FieldGroups` alias to use
`_ListOrTuple`. This is another case where `ModelAdmin.fields` and the
`"fields"` item in fieldsets are presumed by Django to be a `list` or
`tuple`, e.g. there are some system checks that ensure this.

See https://github.com/django/django/blob/afbb8c709d40e77b3f71c152d363c5ad95ceec2d/django/contrib/admin/utils.py#L102-L120

Author: ngnpope

django-stubs/contrib/admin/options.pyi

@@ -56,7 +56,7 @@ class IncorrectLookupParameters(Exception): ...
 FORMFIELD_FOR_DBFIELD_DEFAULTS: Any
 csrf_protect_m: Any
 
-_FieldGroups: TypeAlias = Sequence[str | Sequence[str]]
+_FieldGroups: TypeAlias = _ListOrTuple[str | _ListOrTuple[str]]
 
 @type_check_only
 class _OptionalFieldOpts(TypedDict, total=False):
@@ -67,9 +67,6 @@ class _OptionalFieldOpts(TypedDict, total=False):
 class _FieldOpts(_OptionalFieldOpts, total=True):
     fields: _FieldGroups
 
-# Workaround for mypy issue, a Sequence type should be preferred here.
-# https://github.com/python/mypy/issues/8921
-# _FieldsetSpec = Sequence[Tuple[Optional[str], _FieldOpts]]
 _FieldsetSpec: TypeAlias = _ListOrTuple[tuple[_StrOrPromise | None, _FieldOpts]]
 _ListFilterT: TypeAlias = (
     type[ListFilter]
@@ -82,7 +79,8 @@ _ListFilterT: TypeAlias = (
 # Generic type specifically for models, for use in BaseModelAdmin and subclasses
 # https://github.com/typeddjango/django-stubs/issues/482
 _ModelT = TypeVar("_ModelT", bound=Model)
-_DisplayT: TypeAlias = _ListOrTuple[str | Callable[[_ModelT], str | bool]]
+_DisplayT: TypeAlias = str | Callable[[_ModelT], str | bool]
+_ListDisplayT: TypeAlias = _ListOrTuple[_DisplayT[_ModelT]]
 
 # Options `form`, `list_display`, `list_display_links` and `actions` are not marked as `ClassVar` due to the
 # limitations of the current type system: `ClassVar` cannot contain type variables.
@@ -130,7 +128,7 @@ class BaseModelAdmin(Generic[_ModelT]):
     def get_readonly_fields(self, request: HttpRequest, obj: _ModelT | None = ...) -> _ListOrTuple[str]: ...
     def get_prepopulated_fields(self, request: HttpRequest, obj: _ModelT | None = ...) -> dict[str, Sequence[str]]: ...
     def get_queryset(self, request: HttpRequest) -> QuerySet[_ModelT]: ...
-    def get_sortable_by(self, request: HttpRequest) -> _DisplayT[_ModelT]: ...
+    def get_sortable_by(self, request: HttpRequest) -> _ListDisplayT[_ModelT]: ...
     @overload
     @deprecated("None value for the request parameter will be removed in Django 6.0.")
     def lookup_allowed(self, lookup: str, value: str, request: None = None) -> bool: ...
@@ -150,8 +148,8 @@ _ModelAdmin = TypeVar("_ModelAdmin", bound=ModelAdmin[Any])
 _ActionCallable: TypeAlias = Callable[[_ModelAdmin, HttpRequest, QuerySet[_ModelT]], HttpResponseBase | None]
 
 class ModelAdmin(BaseModelAdmin[_ModelT]):
-    list_display: _DisplayT[_ModelT]
-    list_display_links: _DisplayT[_ModelT] | None
+    list_display: _ListDisplayT[_ModelT]
+    list_display_links: _ListDisplayT[_ModelT] | None
     list_filter: ClassVar[_ListOrTuple[_ListFilterT]]
     list_select_related: ClassVar[bool | _ListOrTuple[str]]
     list_per_page: ClassVar[int]
@@ -220,8 +218,10 @@ class ModelAdmin(BaseModelAdmin[_ModelT]):
         self, request: HttpRequest, default_choices: list[tuple[str, str]] = ...
     ) -> list[tuple[str, str]]: ...
     def get_action(self, action: Callable | str) -> tuple[Callable[..., str], str, str] | None: ...
-    def get_list_display(self, request: HttpRequest) -> _DisplayT[_ModelT]: ...
-    def get_list_display_links(self, request: HttpRequest, list_display: _DisplayT[_ModelT]) -> _DisplayT[_ModelT]: ...
+    def get_list_display(self, request: HttpRequest) -> _ListDisplayT[_ModelT]: ...
+    def get_list_display_links(
+        self, request: HttpRequest, list_display: _ListDisplayT[_ModelT]
+    ) -> _ListDisplayT[_ModelT]: ...
     def get_list_filter(self, request: HttpRequest) -> _ListOrTuple[_ListFilterT]: ...
     def get_list_select_related(self, request: HttpRequest) -> bool | _ListOrTuple[str]: ...
     def get_search_fields(self, request: HttpRequest) -> _ListOrTuple[str]: ...

django-stubs/contrib/admin/utils.pyi

@@ -5,7 +5,7 @@ from typing import Any, Literal, TypeVar, overload, type_check_only
 from uuid import UUID
 
 from _typeshed import Unused
-from django.contrib.admin.options import BaseModelAdmin
+from django.contrib.admin.options import BaseModelAdmin, _DisplayT, _FieldGroups, _FieldsetSpec, _ListDisplayT, _ModelT
 from django.contrib.admin.sites import AdminSite
 from django.db.models.base import Model
 from django.db.models.deletion import Collector
@@ -31,8 +31,11 @@ def prepare_lookup_value(
 def build_q_object_from_lookup_parameters(parameters: dict[str, list[str]]) -> Q: ...
 def quote(s: int | str | UUID) -> str: ...
 def unquote(s: str) -> str: ...
-def flatten(fields: Any) -> list[Callable | str]: ...
-def flatten_fieldsets(fieldsets: Any) -> list[Callable | str]: ...
+@overload
+def flatten(fields: _FieldGroups) -> list[str]: ...
+@overload
+def flatten(fields: _ListDisplayT[_ModelT]) -> list[_DisplayT[_ModelT]]: ...
+def flatten_fieldsets(fieldsets: _FieldsetSpec) -> list[str]: ...
 def get_deleted_objects(
     objs: Sequence[Model | None] | QuerySet[Model], request: HttpRequest, admin_site: AdminSite
 ) -> tuple[list[str], dict[str, int], set[str], list[str]]: ...

django-stubs/contrib/admin/views/main.pyi

@@ -3,7 +3,7 @@ from typing import Any, Literal
 
 from django import forms
 from django.contrib.admin.filters import ListFilter
-from django.contrib.admin.options import ModelAdmin, _DisplayT, _ListFilterT
+from django.contrib.admin.options import ModelAdmin, _ListDisplayT, _ListFilterT
 from django.db.models.base import Model
 from django.db.models.expressions import Expression
 from django.db.models.options import Options
@@ -26,8 +26,8 @@ class ChangeList:
     opts: Options
     lookup_opts: Options
     root_queryset: QuerySet
-    list_display: _DisplayT
-    list_display_links: _DisplayT
+    list_display: _ListDisplayT
+    list_display_links: _ListDisplayT
     list_filter: Sequence[_ListFilterT]
     date_hierarchy: Any
     search_fields: Sequence[str]
@@ -58,8 +58,8 @@ class ChangeList:
         self,
         request: HttpRequest,
         model: type[Model],
-        list_display: _DisplayT,
-        list_display_links: _DisplayT,
+        list_display: _ListDisplayT,
+        list_display_links: _ListDisplayT,
         list_filter: Sequence[_ListFilterT],
         date_hierarchy: str | None,
         search_fields: Sequence[str],

tests/assert_type/contrib/admin/test_utils.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from django import http
+from django.contrib import admin
+from django.contrib.admin.options import _DisplayT
+from django.contrib.admin.utils import flatten, flatten_fieldsets
+from django.db import models
+from typing_extensions import assert_type
+
+
+@admin.display(description="Name")
+def upper_case_name(obj: Person) -> str:
+    return f"{obj.first_name} {obj.last_name}".upper()  # pyright: ignore[reportUnknownMemberType]
+
+
+class Person(models.Model):
+    first_name = models.CharField(max_length=None)  # pyright: ignore[reportUnknownVariableType]
+    last_name = models.CharField(max_length=None)  # pyright: ignore[reportUnknownVariableType]
+    birthday = models.DateField()  # pyright: ignore[reportUnknownVariableType]
+
+
+class PersonListAdmin(admin.ModelAdmin[Person]):
+    fields = [["first_name", "last_name"], "birthday"]
+    list_display = [upper_case_name, "birthday"]
+
+
+class PersonTupleAdmin(admin.ModelAdmin[Person]):
+    fields = (("first_name", "last_name"), "birthday")
+    list_display = (upper_case_name, "birthday")
+
+
+class PersonFieldsetListAdmin(admin.ModelAdmin[Person]):
+    fieldsets = [
+        (
+            "Personal Details",
+            {
+                "description": "Personal details of a person.",
+                "fields": [["first_name", "last_name"], "birthday"],
+            },
+        )
+    ]
+
+
+class PersonFieldsetTupleAdmin(admin.ModelAdmin[Person]):
+    fieldsets = (
+        (
+            "Personal Details",
+            {
+                "description": "Personal details of a person.",
+                "fields": (("first_name", "last_name"), "birthday"),
+            },
+        ),
+    )
+
+
+request = http.HttpRequest()
+admin_site = admin.AdminSite()
+person_list_admin = PersonListAdmin(Person, admin_site)
+person_tuple_admin = PersonTupleAdmin(Person, admin_site)
+person_fieldset_list_admin = PersonFieldsetListAdmin(Person, admin_site)
+person_fieldset_tuple_admin = PersonFieldsetTupleAdmin(Person, admin_site)
+
+# For some reason, pyright cannot see that these are not `None`.
+assert person_list_admin.fields is not None
+assert person_tuple_admin.fields is not None
+assert person_fieldset_list_admin.fieldsets is not None
+assert person_fieldset_tuple_admin.fieldsets is not None
+
+assert_type(flatten(person_list_admin.fields), list[str])
+assert_type(flatten(person_list_admin.get_fields(request)), list[str])
+assert_type(flatten(person_tuple_admin.fields), list[str])
+assert_type(flatten(person_tuple_admin.get_fields(request)), list[str])
+
+assert_type(flatten(person_list_admin.list_display), list[_DisplayT[Person]])
+assert_type(flatten(person_list_admin.get_list_display(request)), list[_DisplayT[Person]])
+assert_type(flatten(person_tuple_admin.list_display), list[_DisplayT[Person]])
+assert_type(flatten(person_tuple_admin.get_list_display(request)), list[_DisplayT[Person]])
+
+assert_type(flatten_fieldsets(person_fieldset_list_admin.fieldsets), list[str])
+assert_type(flatten_fieldsets(person_fieldset_list_admin.get_fieldsets(request)), list[str])
+assert_type(flatten_fieldsets(person_fieldset_tuple_admin.fieldsets), list[str])
+assert_type(flatten_fieldsets(person_fieldset_tuple_admin.get_fieldsets(request)), list[str])