From 51259097fa23ff861eac3f7c63624ea591bf1bd3 Mon Sep 17 00:00:00 2001
From: wyzula-jan <133381102+wyzula-jan@users.noreply.github.com>
Date: Fri, 16 Feb 2024 20:16:19 +0100
Subject: [PATCH] feat(utils.colors): golden_angle_color utility can return
 colors as a  list of QColor, RGB or HEC

---
 bec_widgets/utils/colors.py | 37 ++++++++++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/bec_widgets/utils/colors.py b/bec_widgets/utils/colors.py
index 45de8f4c..f1d65cd4 100644
--- a/bec_widgets/utils/colors.py
+++ b/bec_widgets/utils/colors.py
@@ -1,6 +1,8 @@
+from typing import Literal
+
 import numpy as np
 import pyqtgraph as pg
-from pyqtgraph import mkColor
+from qtpy.QtGui import QColor
 
 
 class Colors:
@@ -10,6 +12,9 @@ class Colors:
 
         Args:
             num (int): Number of angles
+
+        Returns:
+            list: List of angles calculated using the golden ratio.
         """
         phi = 2 * np.pi * ((1 + np.sqrt(5)) / 2)
         angles = []
@@ -21,30 +26,40 @@ class Colors:
         return angles
 
     @staticmethod
-    def golden_angle_color(colormap: str, num: int) -> list:
+    def golden_angle_color(
+        colormap: str, num: int, format: Literal["QColor", "HEX", "RGB"] = "QColor"
+    ) -> list:
         """
-        Extract num colors for from the specified colormap following golden angle distribution.
+        Extract num colors from the specified colormap following golden angle distribution and return them in the specified format.
 
         Args:
-            colormap (str): Name of the colormap
-            num (int): Number of requested colors
+            colormap (str): Name of the colormap.
+            num (int): Number of requested colors.
+            format (Literal["QColor","HEX","RGB"]): The format of the returned colors ('RGB', 'HEX', 'QColor').
 
         Returns:
-            list: List of colors with length <num>
+            list: List of colors in the specified format.
 
         Raises:
             ValueError: If the number of requested colors is greater than the number of colors in the colormap.
         """
-
         cmap = pg.colormap.get(colormap)
-        cmap_colors = cmap.color
+        cmap_colors = cmap.getColors(mode="float")
         if num > len(cmap_colors):
             raise ValueError(
                 f"Number of colors requested ({num}) is greater than the number of colors in the colormap ({len(cmap_colors)})"
             )
         angles = Colors.golden_ratio(len(cmap_colors))
         color_selection = np.round(np.interp(angles, (-np.pi, np.pi), (0, len(cmap_colors))))
-        colors = [
-            mkColor(tuple((cmap_colors[int(ii)] * 255).astype(int))) for ii in color_selection[:num]
-        ]
+        colors = []
+        for ii in color_selection[:num]:
+            color = cmap_colors[int(ii)]
+            if format.upper() == "HEX":
+                colors.append(QColor.fromRgbF(*color).name())
+            elif format.upper() == "RGB":
+                colors.append(tuple((np.array(color) * 255).astype(int)))
+            elif format.upper() == "QCOLOR":
+                colors.append(QColor.fromRgbF(*color))
+            else:
+                raise ValueError("Unsupported format. Please choose 'RGB', 'HEX', or 'QColor'.")
         return colors