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