django/test/utils.py polish (#213)

This PR improves the types in `django/test/utils.py` and fixes `overide_settings.decorate_class` cloaking the type of the decorated class.

See commits for the specific changes.

django-stubs/test/testcases.pyi

@@ -27,7 +27,7 @@ class _AssertTemplateUsedContext:
     template_name: str = ...
     rendered_templates: list[Template] = ...
     rendered_template_names: list[str] = ...
-    context: ContextList = ...
+    context: ContextList[Any] = ...
     def __init__(self, test_case: Any, template_name: Any) -> None: ...
     def on_template_render(
         self, sender: Any, signal: Any, template: Any, context: Any, **kwargs: Any

django-stubs/test/utils.pyi

@@ -1,21 +1,24 @@
 import decimal
 from collections.abc import Callable, Iterable, Iterator, Mapping
-from contextlib import AbstractContextManager, contextmanager
+from contextlib import AbstractContextManager
 from decimal import Decimal
-from io import StringIO
-from typing import Any, TypeVar
+from types import TracebackType
+from typing import Any, TypeVar, overload
+from typing_extensions import Self
 
 from django.apps.registry import Apps
 from django.conf import LazySettings, Settings
 from django.core.checks.registry import CheckRegistry
-from django.db.models.lookups import Lookup, Transform
-from django.db.models.query_utils import RegisterLookupMixin
+from django.db import DefaultConnectionProxy
 from django.test.runner import DiscoverRunner
 from django.test.testcases import SimpleTestCase
 
 _TestClass = type[SimpleTestCase]
 _DecoratedTest = Callable[..., Any] | _TestClass
 _C = TypeVar("_C", bound=Callable[..., Any])
+_T = TypeVar("_T")
+_U = TypeVar("_U")
+_TestClassGeneric = TypeVar("_TestClassGeneric", bound=_TestClass)
 
 TZ_SUPPORT: bool = ...
 
@@ -24,14 +27,25 @@ class Approximate:
     places: int = ...
     def __init__(self, val: Decimal | float, places: int = ...) -> None: ...
 
-class ContextList(list[Any]):
-    def get(self, key: str, default: str | None = ...) -> str: ...
+class ContextList(list[Mapping[str, _T]]):
+    def get(self, key: str, default: _U | None = ...) -> _T | _U | None: ...
     def keys(self) -> set[str]: ...
 
 class _TestState: ...
 
 def setup_test_environment(debug: bool | None = ...) -> None: ...
 def teardown_test_environment() -> None: ...
+def setup_databases(
+    verbosity: int,
+    interactive: bool,
+    *,
+    time_keeper: Any | None = ...,
+    keepdb: bool = ...,
+    debug_sql: bool = ...,
+    parallel: int = ...,
+    aliases: Iterable[str] | None = ...,
+    **kwargs: Any
+) -> list[tuple[DefaultConnectionProxy, str, bool]]: ...
 def get_runner(
     settings: LazySettings, test_runner_class: str | None = ...
 ) -> type[DiscoverRunner]: ...
@@ -42,39 +56,47 @@ class TestContextDecorator:
     def __init__(
         self, attr_name: str | None = ..., kwarg_name: str | None = ...
     ) -> None: ...
-    def enable(self) -> Any: ...
+    def enable(self) -> Any | None: ...
     def disable(self) -> None: ...
-    def __enter__(self) -> Apps | None: ...
-    def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None: ...
-    def decorate_class(self, cls: _TestClass) -> _TestClass: ...
+    def __enter__(self) -> Any | None: ...
+    def __exit__(
+        self,
+        exc_type: type[Exception] | None,
+        exc_value: Exception | None,
+        traceback: TracebackType | None,
+    ) -> None: ...
+    def decorate_class(self, cls: _TestClassGeneric) -> _TestClassGeneric: ...
     def decorate_callable(self, func: _C) -> _C: ...
-    def __call__(self, decorated: _DecoratedTest) -> Any: ...
+    @overload
+    def __call__(self, decorated: _TestClassGeneric) -> _TestClassGeneric: ...
+    @overload
+    def __call__(self, decorated: _C) -> _C: ...
 
 class override_settings(TestContextDecorator):
+    enable_exception: bool | None = ...
+    wrapped: Settings = ...
     options: dict[str, Any] = ...
     def __init__(self, **kwargs: Any) -> None: ...
-    wrapped: Settings = ...
     def save_options(self, test_func: _DecoratedTest) -> None: ...
-    def decorate_class(self, cls: type) -> type: ...
 
 class modify_settings(override_settings):
     wrapped: Settings
     operations: list[tuple[str, dict[str, list[str] | str]]] = ...
+    options: dict[str, list[tuple[str, str] | str]] = ...
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     def save_options(self, test_func: _DecoratedTest) -> None: ...
-    options: dict[str, list[tuple[str, str] | str]] = ...
 
 class override_system_checks(TestContextDecorator):
     registry: CheckRegistry = ...
     new_checks: list[Callable[..., Any]] = ...
     deployment_checks: list[Callable[..., Any]] | None = ...
+    old_checks: set[Callable[..., Any]] = ...
+    old_deployment_checks: set[Callable[..., Any]] = ...
     def __init__(
         self,
         new_checks: list[Callable[..., Any]],
         deployment_checks: list[Callable[..., Any]] | None = ...,
     ) -> None: ...
-    old_checks: set[Callable[..., Any]] = ...
-    old_deployment_checks: set[Callable[..., Any]] = ...
 
 class CaptureQueriesContext:
     connection: Any = ...
@@ -87,7 +109,7 @@ class CaptureQueriesContext:
     def __len__(self) -> int: ...
     @property
     def captured_queries(self) -> list[dict[str, str]]: ...
-    def __enter__(self) -> CaptureQueriesContext: ...
+    def __enter__(self) -> Self: ...
     def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None: ...
 
 class ignore_warnings(TestContextDecorator):
@@ -117,40 +139,4 @@ class isolate_apps(TestContextDecorator):
     def __init__(self, *installed_apps: Any, **kwargs: Any) -> None: ...
     old_apps: Apps = ...
 
-@contextmanager
-def extend_sys_path(*paths: str) -> Iterator[None]: ...
-@contextmanager
-def captured_output(stream_name: Any) -> Iterator[StringIO]: ...
-@contextmanager
-def captured_stdin() -> Iterator[StringIO]: ...
-@contextmanager
-def captured_stdout() -> Iterator[StringIO]: ...
-@contextmanager
-def captured_stderr() -> Iterator[StringIO]: ...
-@contextmanager
-def freeze_time(t: float) -> Iterator[None]: ...
-def tag(*tags: str) -> Any: ...
-
-_Signature = str
-_TestDatabase = tuple[str, list[str]]
-
-def dependency_ordered(
-    test_databases: Iterable[tuple[_Signature, _TestDatabase]],
-    dependencies: Mapping[str, list[str]],
-) -> list[tuple[_Signature, _TestDatabase]]: ...
-def get_unique_databases_and_mirrors() -> (
-    tuple[dict[_Signature, _TestDatabase], dict[str, Any]]
-): ...
-def teardown_databases(
-    old_config: Iterable[tuple[Any, str, bool]],
-    verbosity: int,
-    parallel: int = ...,
-    keepdb: bool = ...,
-) -> None: ...
-def require_jinja2(test_func: _C) -> _C: ...
-@contextmanager
-def register_lookup(
-    field: type[RegisterLookupMixin],
-    *lookups: type[Lookup[Any] | Transform],
-    lookup_name: str | None = ...
-) -> Iterator[None]: ...
+def tag(*tags: str) -> Callable[[_T], _T]: ...