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 (
8 TYPE_CHECKING,
9 Any,
10)
11import warnings
13if TYPE_CHECKING:
14 from collections.abc import Generator
15 from types import FrameType
18@contextlib.contextmanager
19def rewrite_exception(old_name: str, new_name: str) -> Generator[None]:
20 """
21 Rewrite the message of an exception.
22 """
23 try:
24 yield
25 except Exception as err:
26 if not err.args:
27 raise
28 msg = str(err.args[0])
29 msg = msg.replace(old_name, new_name)
30 args: tuple[Any, ...] = (msg,)
31 if len(err.args) > 1:
32 args = args + err.args[1:]
33 err.args = args
34 raise
37def find_stack_level() -> int:
38 """
39 Find the first place in the stack that is not inside pandas
40 (tests notwithstanding).
41 """
43 import pandas as pd
45 pkg_dir = os.path.dirname(pd.__file__)
46 test_dir = os.path.join(pkg_dir, "tests")
48 # https://stackoverflow.com/questions/17407119/python-inspect-stack-is-slow
49 frame: FrameType | None = inspect.currentframe()
50 try:
51 n = 0
52 while frame:
53 filename = inspect.getfile(frame)
54 if filename.startswith(pkg_dir) and not filename.startswith(test_dir):
55 frame = frame.f_back
56 n += 1
57 else:
58 break
59 finally:
60 # See note in
61 # https://docs.python.org/3/library/inspect.html#inspect.Traceback
62 del frame
63 return n
66@contextlib.contextmanager
67def rewrite_warning(
68 target_message: str,
69 target_category: type[Warning],
70 new_message: str,
71 new_category: type[Warning] | None = None,
72) -> Generator[None]:
73 """
74 Rewrite the message of a warning.
76 Parameters
77 ----------
78 target_message : str
79 Warning message to match.
80 target_category : Warning
81 Warning type to match.
82 new_message : str
83 New warning message to emit.
84 new_category : Warning or None, default None
85 New warning type to emit. When None, will be the same as target_category.
86 """
87 if new_category is None:
88 new_category = target_category
89 with warnings.catch_warnings(record=True) as record:
90 yield
91 if len(record) > 0:
92 match = re.compile(target_message)
93 for warning in record:
94 if warning.category is target_category and re.search(
95 match, str(warning.message)
96 ):
97 category = new_category
98 message: Warning | str = new_message
99 else:
100 category, message = warning.category, warning.message
101 warnings.warn_explicit(
102 message=message,
103 category=category,
104 filename=warning.filename,
105 lineno=warning.lineno,
106 )