From 657c21ea911f4579cf4117328ce23afc5e49fe68 Mon Sep 17 00:00:00 2001
From: AlexWaygood <alex.waygood@gmail.com>
Date: Sat, 2 Jul 2022 15:38:07 +0100
Subject: [PATCH] Improve `warnings.catch_warnings`

---
 stdlib/warnings.pyi                      | 41 ++++++++++++------------
 tests/stubtest_allowlists/py3_common.txt |  1 -
 2 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/stdlib/warnings.pyi b/stdlib/warnings.pyi
index c9c143991e3a..5cc6b946409b 100644
--- a/stdlib/warnings.pyi
+++ b/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,11 +57,11 @@ 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 = ...,
@@ -68,10 +69,10 @@ class catch_warnings:
             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 = ...,
@@ -79,10 +80,10 @@ class catch_warnings:
             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 = ...,
@@ -90,22 +91,20 @@ class catch_warnings:
             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]: ...
diff --git a/tests/stubtest_allowlists/py3_common.txt b/tests/stubtest_allowlists/py3_common.txt
index 6a9d5d1ef420..85cb55ed3cc1 100644
--- a/tests/stubtest_allowlists/py3_common.txt
+++ b/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