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
37_pkg_dir: str | None = None
38_test_dir: str | None = None
41def find_stack_level() -> int:
42 """
43 Find the first place in the stack that is not inside pandas
44 (tests notwithstanding).
45 """
46 global _pkg_dir, _test_dir
47 if _pkg_dir is None or _test_dir is None:
48 import pandas as pd
50 _pkg_dir = os.path.dirname(pd.__file__) + os.sep
51 _test_dir = os.path.join(_pkg_dir, "tests") + os.sep
53 pkg_dir = _pkg_dir
54 test_dir = _test_dir
55 assert pkg_dir is not None
56 assert test_dir is not None
58 # https://stackoverflow.com/questions/17407119/python-inspect-stack-is-slow
59 frame: FrameType | None = inspect.currentframe()
60 try:
61 n = 0
62 while frame:
63 filename = frame.f_code.co_filename
64 if filename.startswith(pkg_dir) and not filename.startswith(test_dir):
65 frame = frame.f_back
66 n += 1
67 else:
68 break
69 finally:
70 # See note in
71 # https://docs.python.org/3/library/inspect.html#inspect.Traceback
72 del frame
73 return n
76@contextlib.contextmanager
77def rewrite_warning(
78 target_message: str,
79 target_category: type[Warning],
80 new_message: str,
81 new_category: type[Warning] | None = None,
82) -> Generator[None]:
83 """
84 Rewrite the message of a warning.
86 Parameters
87 ----------
88 target_message : str
89 Warning message to match.
90 target_category : Warning
91 Warning type to match.
92 new_message : str
93 New warning message to emit.
94 new_category : Warning or None, default None
95 New warning type to emit. When None, will be the same as target_category.
96 """
97 if new_category is None:
98 new_category = target_category
99 with warnings.catch_warnings(record=True) as record:
100 yield
101 if len(record) > 0:
102 match = re.compile(target_message)
103 for warning in record:
104 if warning.category is target_category and re.search(
105 match, str(warning.message)
106 ):
107 category = new_category
108 message: Warning | str = new_message
109 else:
110 category, message = warning.category, warning.message
111 warnings.warn_explicit(
112 message=message,
113 category=category,
114 filename=warning.filename,
115 lineno=warning.lineno,
116 )