From 41bf6a19822d6694973449d795f8bfe1d50d5a03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20S=C5=82apek?=
 <28485371+mslapek@users.noreply.github.com>
Date: Tue, 26 Nov 2019 04:35:27 +0100
Subject: [PATCH] Refactor warnings.catch_warning to be a class. (#3499)

---
 stdlib/2and3/warnings.pyi | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/stdlib/2and3/warnings.pyi b/stdlib/2and3/warnings.pyi
index 2e95533f7..6e2c3a53f 100644
--- a/stdlib/2and3/warnings.pyi
+++ b/stdlib/2and3/warnings.pyi
@@ -1,8 +1,8 @@
 # Stubs for warnings
 
 import sys
-from typing import Any, Dict, List, NamedTuple, Optional, overload, TextIO, Tuple, Type, Union, ContextManager
-from types import ModuleType
+from typing import Any, Dict, List, NamedTuple, Optional, overload, TextIO, Tuple, Type, Union
+from types import ModuleType, TracebackType
 
 if sys.version_info >= (3, 8):
     from typing import Literal
@@ -43,12 +43,20 @@ class _Record(NamedTuple):
     file: Optional[TextIO]
     line: Optional[str]
 
+class catch_warnings:
+    @overload
+    def __new__(cls, *, record: Literal[False] = ..., module: Optional[ModuleType] = ...) -> _catch_warnings_without_records: ...
+    @overload
+    def __new__(cls, *, record: Literal[True], module: Optional[ModuleType] = ...) -> _catch_warnings_with_records: ...
+    @overload
+    def __new__(cls, *, record: bool, module: Optional[ModuleType] = ...) -> catch_warnings: ...
+    def __enter__(self) -> Optional[List[_Record]]: ...
+    def __exit__(self, exc_type: Optional[Type[BaseException]],
+                 exc_val: Optional[BaseException],
+                 exc_tb: Optional[TracebackType]) -> None: ...
 
-@overload
-def catch_warnings(*, record: Literal[False] = ..., module: Optional[ModuleType] = ...) -> ContextManager[None]: ...
+class _catch_warnings_without_records(catch_warnings):
+    def __enter__(self) -> None: ...
 
-@overload
-def catch_warnings(*, record: Literal[True], module: Optional[ModuleType] = ...) -> ContextManager[List[_Record]]: ...
-
-@overload
-def catch_warnings(*, record: bool, module: Optional[ModuleType] = ...) -> ContextManager[Optional[List[_Record]]]: ...
+class _catch_warnings_with_records(catch_warnings):
+    def __enter__(self) -> List[_Record]: ...