warnings_utils.py
1 import warnings 2 3 # ANSI escape code 4 ANSI_BASE = "\033[" 5 COLORS = { 6 "default_bold": f"{ANSI_BASE}1m", 7 "red": f"{ANSI_BASE}31m", 8 "red_bold": f"{ANSI_BASE}1;31m", 9 "yellow": f"{ANSI_BASE}33m", 10 "yellow_bold": f"{ANSI_BASE}1;33m", 11 "blue": f"{ANSI_BASE}34m", 12 "blue_bold": f"{ANSI_BASE}1;34m", 13 } 14 RESET = "\033[0m" 15 16 17 def color_warning(message: str, stacklevel: int, color: str, category: type[Warning] = UserWarning): 18 if color in COLORS: 19 message = f"{COLORS[color]}{message}{RESET}" 20 21 warnings.warn( 22 message=message, 23 category=category, 24 stacklevel=stacklevel + 1, 25 )