1from __future__ import annotations
2
3from contextlib import (
4 AbstractContextManager,
5 contextmanager,
6 nullcontext,
7)
8import inspect
9import re
10import sys
11from typing import (
12 TYPE_CHECKING,
13 Literal,
14 Union,
15 cast,
16)
17import warnings
18
19if TYPE_CHECKING:
20 from collections.abc import (
21 Generator,
22 Sequence,
23 )
24
25
26@contextmanager
27def assert_produces_warning(
28 expected_warning: type[Warning] | bool | tuple[type[Warning], ...] | None = Warning,
29 filter_level: Literal[
30 "error", "ignore", "always", "default", "module", "once"
31 ] = "always",
32 check_stacklevel: bool = True,
33 raise_on_extra_warnings: bool = True,
34 match: str | tuple[str | None, ...] | None = None,
35 must_find_all_warnings: bool = True,
36) -> Generator[list[warnings.WarningMessage]]:
37 """
38 Context manager for running code expected to either raise a specific warning,
39 multiple specific warnings, or not raise any warnings. Verifies that the code
40 raises the expected warning(s), and that it does not raise any other unexpected
41 warnings. It is basically a wrapper around ``warnings.catch_warnings``.
42
43 Parameters
44 ----------
45 expected_warning : {Warning, False, tuple[Warning, ...], None}, default Warning
46 The type of Exception raised. ``exception.Warning`` is the base
47 class for all warnings. To raise multiple types of exceptions,
48 pass them as a tuple. To check that no warning is returned,
49 specify ``False`` or ``None``.
50 filter_level : str or None, default "always"
51 Specifies whether warnings are ignored, displayed, or turned
52 into errors.
53 Valid values are:
54
55 * "error" - turns matching warnings into exceptions
56 * "ignore" - discard the warning
57 * "always" - always emit a warning
58 * "default" - print the warning the first time it is generated
59 from each location
60 * "module" - print the warning the first time it is generated
61 from each module
62 * "once" - print the warning the first time it is generated
63
64 check_stacklevel : bool, default True
65 If True, displays the line that called the function containing
66 the warning to show were the function is called. Otherwise, the
67 line that implements the function is displayed.
68 raise_on_extra_warnings : bool, default True
69 Whether extra warnings not of the type `expected_warning` should
70 cause the test to fail.
71 match : {str, tuple[str, ...]}, optional
72 Match warning message. If it's a tuple, it has to be the size of
73 `expected_warning`. If additionally `must_find_all_warnings` is
74 True, each expected warning's message gets matched with a respective
75 match. Otherwise, multiple values get treated as an alternative.
76 must_find_all_warnings : bool, default True
77 If True and `expected_warning` is a tuple, each expected warning
78 type must get encountered. Otherwise, even one expected warning
79 results in success.
80
81 Examples
82 --------
83 >>> import warnings
84 >>> with assert_produces_warning():
85 ... warnings.warn(UserWarning())
86 >>> with assert_produces_warning(False):
87 ... warnings.warn(RuntimeWarning())
88 Traceback (most recent call last):
89 ...
90 AssertionError: Caused unexpected warning(s): ['RuntimeWarning'].
91 >>> with assert_produces_warning(UserWarning):
92 ... warnings.warn(RuntimeWarning())
93 Traceback (most recent call last):
94 ...
95 AssertionError: Did not see expected warning of class 'UserWarning'.
96
97 ..warn:: This is *not* thread-safe.
98 """
99 __tracebackhide__ = True
100
101 with warnings.catch_warnings(record=True) as w:
102 warnings.simplefilter(filter_level)
103 try:
104 yield w
105 finally:
106 if expected_warning:
107 if isinstance(expected_warning, tuple) and must_find_all_warnings:
108 match = (
109 match
110 if isinstance(match, tuple)
111 else (match,) * len(expected_warning)
112 )
113 for warning_type, warning_match in zip(
114 expected_warning, match, strict=True
115 ):
116 _assert_caught_expected_warnings(
117 caught_warnings=w,
118 expected_warning=warning_type,
119 match=warning_match,
120 check_stacklevel=check_stacklevel,
121 )
122 else:
123 expected_warning = cast(
124 Union[type[Warning], tuple[type[Warning], ...]],
125 expected_warning,
126 )
127 match = (
128 "|".join(m for m in match if m)
129 if isinstance(match, tuple)
130 else match
131 )
132 _assert_caught_expected_warnings(
133 caught_warnings=w,
134 expected_warning=expected_warning,
135 match=match,
136 check_stacklevel=check_stacklevel,
137 )
138 if raise_on_extra_warnings:
139 _assert_caught_no_extra_warnings(
140 caught_warnings=w,
141 expected_warning=expected_warning,
142 )
143
144
145def maybe_produces_warning(
146 warning: type[Warning], condition: bool, **kwargs
147) -> AbstractContextManager:
148 """
149 Return a context manager that possibly checks a warning based on the condition
150 """
151 if condition:
152 return assert_produces_warning(warning, **kwargs)
153 else:
154 return nullcontext()
155
156
157def _assert_caught_expected_warnings(
158 *,
159 caught_warnings: Sequence[warnings.WarningMessage],
160 expected_warning: type[Warning] | tuple[type[Warning], ...],
161 match: str | None,
162 check_stacklevel: bool,
163) -> None:
164 """Assert that there was the expected warning among the caught warnings."""
165 saw_warning = False
166 matched_message = False
167 unmatched_messages = []
168 warning_name = (
169 tuple(x.__name__ for x in expected_warning)
170 if isinstance(expected_warning, tuple)
171 else expected_warning.__name__
172 )
173
174 for actual_warning in caught_warnings:
175 if issubclass(actual_warning.category, expected_warning):
176 saw_warning = True
177
178 if check_stacklevel:
179 _assert_raised_with_correct_stacklevel(actual_warning)
180
181 if match is not None:
182 if re.search(match, str(actual_warning.message)):
183 matched_message = True
184 else:
185 unmatched_messages.append(actual_warning.message)
186
187 if not saw_warning:
188 raise AssertionError(f"Did not see expected warning of class {warning_name!r}")
189
190 if match and not matched_message:
191 raise AssertionError(
192 f"Did not see warning {warning_name!r} "
193 f"matching '{match}'. The emitted warning messages are "
194 f"{unmatched_messages}"
195 )
196
197
198def _assert_caught_no_extra_warnings(
199 *,
200 caught_warnings: Sequence[warnings.WarningMessage],
201 expected_warning: type[Warning] | bool | tuple[type[Warning], ...] | None,
202) -> None:
203 """Assert that no extra warnings apart from the expected ones are caught."""
204 extra_warnings = []
205
206 for actual_warning in caught_warnings:
207 if _is_unexpected_warning(actual_warning, expected_warning):
208 # GH#38630 pytest.filterwarnings does not suppress these.
209 if actual_warning.category == ResourceWarning:
210 # GH 44732: Don't make the CI flaky by filtering SSL-related
211 # ResourceWarning from dependencies
212 if "unclosed <ssl.SSLSocket" in str(actual_warning.message):
213 continue
214 # GH 44844: Matplotlib leaves font files open during the entire process
215 # upon import. Don't make CI flaky if ResourceWarning raised
216 # due to these open files.
217 if any("matplotlib" in mod for mod in sys.modules):
218 continue
219 if actual_warning.category == EncodingWarning:
220 # EncodingWarnings are checked in the CI
221 # pyproject.toml errors on EncodingWarnings in pandas
222 # Ignore EncodingWarnings from other libraries
223 continue
224 extra_warnings.append(
225 (
226 actual_warning.category.__name__,
227 actual_warning.message,
228 actual_warning.filename,
229 actual_warning.lineno,
230 )
231 )
232
233 if extra_warnings:
234 raise AssertionError(f"Caused unexpected warning(s): {extra_warnings!r}")
235
236
237def _is_unexpected_warning(
238 actual_warning: warnings.WarningMessage,
239 expected_warning: type[Warning] | bool | tuple[type[Warning], ...] | None,
240) -> bool:
241 """Check if the actual warning issued is unexpected."""
242 if actual_warning and not expected_warning:
243 return True
244 expected_warning = cast(type[Warning], expected_warning)
245 return bool(not issubclass(actual_warning.category, expected_warning))
246
247
248def _assert_raised_with_correct_stacklevel(
249 actual_warning: warnings.WarningMessage,
250) -> None:
251 # https://stackoverflow.com/questions/17407119/python-inspect-stack-is-slow
252 frame = inspect.currentframe()
253 for _ in range(4):
254 frame = frame.f_back # type: ignore[union-attr]
255 try:
256 caller_filename = inspect.getfile(frame) # type: ignore[arg-type]
257 finally:
258 # See note in
259 # https://docs.python.org/3/library/inspect.html#inspect.Traceback
260 del frame
261 msg = (
262 "Warning not set with correct stacklevel. "
263 f"File where warning is raised: {actual_warning.filename} != "
264 f"{caller_filename}. Warning message: {actual_warning.message}"
265 )
266 assert actual_warning.filename == caller_filename, msg