_exceptions.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. from __future__ import annotations
  2. import contextlib
  3. import inspect
  4. import os
  5. import re
  6. from typing import TYPE_CHECKING
  7. import warnings
  8. if TYPE_CHECKING:
  9. from collections.abc import Generator
  10. from types import FrameType
  11. @contextlib.contextmanager
  12. def 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
  28. def find_stack_level() -> int:
  29. """
  30. Find the first place in the stack that is not inside pandas
  31. (tests notwithstanding).
  32. """
  33. import pandas as pd
  34. pkg_dir = os.path.dirname(pd.__file__)
  35. test_dir = os.path.join(pkg_dir, "tests")
  36. # https://stackoverflow.com/questions/17407119/python-inspect-stack-is-slow
  37. frame: FrameType | None = inspect.currentframe()
  38. try:
  39. n = 0
  40. while frame:
  41. filename = inspect.getfile(frame)
  42. if filename.startswith(pkg_dir) and not filename.startswith(test_dir):
  43. frame = frame.f_back
  44. n += 1
  45. else:
  46. break
  47. finally:
  48. # See note in
  49. # https://docs.python.org/3/library/inspect.html#inspect.Traceback
  50. del frame
  51. return n
  52. @contextlib.contextmanager
  53. def rewrite_warning(
  54. target_message: str,
  55. target_category: type[Warning],
  56. new_message: str,
  57. new_category: type[Warning] | None = None,
  58. ) -> Generator[None, None, None]:
  59. """
  60. Rewrite the message of a warning.
  61. Parameters
  62. ----------
  63. target_message : str
  64. Warning message to match.
  65. target_category : Warning
  66. Warning type to match.
  67. new_message : str
  68. New warning message to emit.
  69. new_category : Warning or None, default None
  70. New warning type to emit. When None, will be the same as target_category.
  71. """
  72. if new_category is None:
  73. new_category = target_category
  74. with warnings.catch_warnings(record=True) as record:
  75. yield
  76. if len(record) > 0:
  77. match = re.compile(target_message)
  78. for warning in record:
  79. if warning.category is target_category and re.search(
  80. match, str(warning.message)
  81. ):
  82. category = new_category
  83. message: Warning | str = new_message
  84. else:
  85. category, message = warning.category, warning.message
  86. warnings.warn_explicit(
  87. message=message,
  88. category=category,
  89. filename=warning.filename,
  90. lineno=warning.lineno,
  91. )