/ mlflow / utils / warnings_utils.py
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      )