Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import contextlib
4import inspect
5import os
6import re
7from typing import Generator
8import warnings
11@contextlib.contextmanager
12def rewrite_exception(old_name: str, new_name: str) -> Generator[None, None, None]:
13 """
14 Rewrite the message of an exception.
15 """
16 try:
17 yield
18 except Exception as err:
19 if not err.args:
20 raise
21 msg = str(err.args[0])
22 msg = msg.replace(old_name, new_name)
23 args: tuple[str, ...] = (msg,)
24 if len(err.args) > 1:
25 args = args + err.args[1:]
26 err.args = args
27 raise
30def find_stack_level() -> int:
31 """
32 Find the first place in the stack that is not inside pandas
33 (tests notwithstanding).
34 """
36 import pandas as pd
38 pkg_dir = os.path.dirname(pd.__file__)
39 test_dir = os.path.join(pkg_dir, "tests")
41 # https://stackoverflow.com/questions/17407119/python-inspect-stack-is-slow
42 frame = inspect.currentframe()
43 n = 0
44 while frame:
45 fname = inspect.getfile(frame)
46 if fname.startswith(pkg_dir) and not fname.startswith(test_dir):
47 frame = frame.f_back
48 n += 1
49 else:
50 break
51 return n
54@contextlib.contextmanager
55def rewrite_warning(
56 target_message: str,
57 target_category: type[Warning],
58 new_message: str,
59 new_category: type[Warning] | None = None,
60) -> Generator[None, None, None]:
61 """
62 Rewrite the message of a warning.
64 Parameters
65 ----------
66 target_message : str
67 Warning message to match.
68 target_category : Warning
69 Warning type to match.
70 new_message : str
71 New warning message to emit.
72 new_category : Warning or None, default None
73 New warning type to emit. When None, will be the same as target_category.
74 """
75 if new_category is None:
76 new_category = target_category
77 with warnings.catch_warnings(record=True) as record:
78 yield
79 if len(record) > 0:
80 match = re.compile(target_message)
81 for warning in record:
82 if warning.category is target_category and re.search(
83 match, str(warning.message)
84 ):
85 category = new_category
86 message: Warning | str = new_message
87 else:
88 category, message = warning.category, warning.message
89 warnings.warn_explicit(
90 message=message,
91 category=category,
92 filename=warning.filename,
93 lineno=warning.lineno,
94 )