Improve type hints for test client response.context (#1100)

* Improve type hints for test client response.context

* fix test

django-stubs/test/client.pyi

@@ -28,6 +28,7 @@ from django.http.cookie import SimpleCookie
 from django.http.request import HttpRequest
 from django.http.response import HttpResponseBase
 from django.template.base import Template
+from django.test.utils import ContextList
 from django.urls import ResolverMatch
 
 BOUNDARY: str = ...
@@ -125,7 +126,7 @@ class _MonkeyPatchedWSGIResponse(_WSGIResponse):
     request: Dict[str, Any]
     client: Client
     templates: List[Template]
-    context: List[Dict[str, Any]]
+    context: ContextList | Dict[str, Any]
     content: bytes
     resolver_match: ResolverMatch
 
@@ -134,7 +135,7 @@ class _MonkeyPatchedASGIResponse(_ASGIResponse):
     request: Dict[str, Any]
     client: AsyncClient
     templates: List[Template]
-    context: List[Dict[str, Any]]
+    context: ContextList | Dict[str, Any]
     content: bytes
     resolver_match: ResolverMatch
 

django-stubs/test/utils.pyi

@@ -29,6 +29,7 @@ from django.db.models.lookups import Lookup, Transform
 from django.db.models.query_utils import RegisterLookupMixin
 from django.test.runner import DiscoverRunner
 from django.test.testcases import SimpleTestCase
+from typing_extensions import SupportsIndex
 
 _TestClass = Type[SimpleTestCase]
 _DecoratedTest = Union[Callable, _TestClass]
@@ -41,8 +42,10 @@ class Approximate:
     places: int = ...
     def __init__(self, val: Union[Decimal, float], places: int = ...) -> None: ...
 
-class ContextList(list):
+class ContextList(List[Dict[str, Any]]):
-    def get(self, key: str, default: Optional[str] = ...) -> str: ...
+    def __getitem__(self, key: Union[str, SupportsIndex, slice]) -> Any: ...
+    def get(self, key: str, default: Optional[Any] = ...) -> Any: ...
+    def __contains__(self, key: object) -> bool: ...
     def keys(self) -> Set[str]: ...
 
 class _TestState: ...

tests/typecheck/test/test_client.yml

@@ -7,7 +7,7 @@
       reveal_type(response.request)  # N: Revealed type is "builtins.dict[builtins.str, Any]"
       reveal_type(response.templates)  # N: Revealed type is "builtins.list[django.template.base.Template]"
       reveal_type(response.client)  # N: Revealed type is "django.test.client.Client"
-      reveal_type(response.context)  # N: Revealed type is "builtins.list[builtins.dict[builtins.str, Any]]"
+      reveal_type(response.context)  # N: Revealed type is "Union[django.test.utils.ContextList, builtins.dict[builtins.str, Any]]"
       reveal_type(response.content)  # N: Revealed type is "builtins.bytes"
       response.json()
 -   case: async_client_methods
@@ -20,7 +20,7 @@
         reveal_type(response.request)  # N: Revealed type is "builtins.dict[builtins.str, Any]"
         reveal_type(response.templates)  # N: Revealed type is "builtins.list[django.template.base.Template]"
         reveal_type(response.client)  # N: Revealed type is "django.test.client.AsyncClient"
-        reveal_type(response.context)  # N: Revealed type is "builtins.list[builtins.dict[builtins.str, Any]]"
+        reveal_type(response.context)  # N: Revealed type is "Union[django.test.utils.ContextList, builtins.dict[builtins.str, Any]]"
         reveal_type(response.content)  # N: Revealed type is "builtins.bytes"
         response.json()
 -   case: request_factories