Improve warnings.catch_warnings

stdlib/warnings.pyi

@@ -2,7 +2,7 @@ import sys
 from _warnings import warn as warn, warn_explicit as warn_explicit
 from collections.abc import Sequence
 from types import ModuleType, TracebackType
-from typing import Any, TextIO, overload
+from typing import Any, Generic, TextIO, TypeVar, overload
 from typing_extensions import Literal, TypeAlias
 
 __all__ = [
@@ -16,6 +16,7 @@ __all__ = [
     "catch_warnings",
 ]
 
+_W = TypeVar("_W", bound=list[WarningMessage] | None)
 _ActionKind: TypeAlias = Literal["default", "error", "ignore", "always", "module", "once"]
 
 filters: Sequence[tuple[str, str | None, type[Warning], str | None, int]]  # undocumented, do not mutate
@@ -56,56 +57,54 @@ class WarningMessage:
         source: Any | None = ...,
     ) -> None: ...
 
-class catch_warnings:
+class catch_warnings(Generic[_W]):
     if sys.version_info >= (3, 11):
         @overload
-        def __new__(
-            cls,
+        def __init__(
+            self: catch_warnings[None],
             *,
             record: Literal[False] = ...,
             module: ModuleType | None = ...,
             action: _ActionKind | None = ...,
             category: type[Warning] = ...,
             lineno: int = ...,
             append: bool = ...,
-        ) -> _catch_warnings_without_records: ...
+        ) -> None: ...
         @overload
-        def __new__(
-            cls,
+        def __init__(
+            self: catch_warnings[list[WarningMessage]],
             *,
             record: Literal[True],
             module: ModuleType | None = ...,
             action: _ActionKind | None = ...,
             category: type[Warning] = ...,
             lineno: int = ...,
             append: bool = ...,
-        ) -> _catch_warnings_with_records: ...
+        ) -> None: ...
         @overload
-        def __new__(
-            cls,
+        def __init__(
+            self: catch_warnings[list[WarningMessage] | None],
             *,
             record: bool,
             module: ModuleType | None = ...,
             action: _ActionKind | None = ...,
             category: type[Warning] = ...,
             lineno: int = ...,
             append: bool = ...,
-        ) -> catch_warnings: ...
+        ) -> None: ...
     else:
         @overload
-        def __new__(cls, *, record: Literal[False] = ..., module: ModuleType | None = ...) -> _catch_warnings_without_records: ...
+        def __init__(self: catch_warnings[None], *, record: Literal[False] = ..., module: ModuleType | None = ...) -> None: ...
         @overload
-        def __new__(cls, *, record: Literal[True], module: ModuleType | None = ...) -> _catch_warnings_with_records: ...
+        def __init__(
+            self: catch_warnings[list[WarningMessage]], *, record: Literal[True], module: ModuleType | None = ...
+        ) -> None: ...
         @overload
-        def __new__(cls, *, record: bool, module: ModuleType | None = ...) -> catch_warnings: ...
+        def __init__(
+            self: catch_warnings[list[WarningMessage] | None], *, record: bool, module: ModuleType | None = ...
+        ) -> None: ...
 
-    def __enter__(self) -> list[WarningMessage] | None: ...
+    def __enter__(self) -> _W: ...
     def __exit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None: ...
-
-class _catch_warnings_without_records(catch_warnings):
-    def __enter__(self) -> None: ...
-
-class _catch_warnings_with_records(catch_warnings):
-    def __enter__(self) -> list[WarningMessage]: ...

tests/stubtest_allowlists/py3_common.txt

@@ -216,7 +216,6 @@ typing.type_check_only  # typing decorator that is not available at runtime
 unittest.mock.patch  # It's a complicated overload and I haven't been able to figure out why stubtest doesn't like it
 urllib.parse._DefragResultBase.__new__  # Generic NamedTuple is problematic in mypy, so regular tuple was used. See https://github.com/python/mypy/issues/685
 urllib.request.HTTPPasswordMgrWithPriorAuth.__init__  # Args are passed as is to super, so super args are specified
-warnings.catch_warnings.__init__  # Defining this ruins the __new__ overrides
 weakref.CallableProxyType.__getattr__  # Should have all attributes of proxy
 weakref.ProxyType.__getattr__  # Should have all attributes of proxy
 weakref.ReferenceType.*  # Alias for _weakref.ReferenceType, problems should be fixed there