star

venv/lib/python3.9/site-packages/mlflow/utils/warnings_utils.py

@@ -0,0 +1,25 @@
+import warnings
+
+# ANSI escape code
+ANSI_BASE = "\033["
+COLORS = {
+    "default_bold": f"{ANSI_BASE}1m",
+    "red": f"{ANSI_BASE}31m",
+    "red_bold": f"{ANSI_BASE}1;31m",
+    "yellow": f"{ANSI_BASE}33m",
+    "yellow_bold": f"{ANSI_BASE}1;33m",
+    "blue": f"{ANSI_BASE}34m",
+    "blue_bold": f"{ANSI_BASE}1;34m",
+}
+RESET = "\033[0m"
+
+
+def color_warning(message: str, stacklevel: int, color: str, category: type[Warning] = UserWarning):
+    if color in COLORS:
+        message = f"{COLORS[color]}{message}{RESET}"
+
+    warnings.warn(
+        message=message,
+        category=category,
+        stacklevel=stacklevel + 1,
+    )