Merge pull request #1853 from NilsNemitz/getColor-getStops

make getColors and getStops behavior more consistent; add tests

pyqtgraph/colormap.py

@@ -5,6 +5,8 @@ from os import path, listdir
 from collections.abc import Callable, Sequence
 import warnings
 
+__all__ = ['ColorMap']
+
 _mapCache = {}
 
 def listMaps(source=None):
@@ -598,32 +600,45 @@ class ColorMap(object):
         pen.setCosmetic(True)
         return pen
 
-    def getColors(self, mode=None):
-        """Returns a list of all color stops, converted to the specified mode.
-        If `mode` is None, no conversion is performed.
+    def getColors(self, mode=BYTE):
+        """
+        Returns a list of the colors associated with the stops of the color map.
+        
+        The parameter `mode` can be one of
+            - `ColorMap.BYTE` or 'byte' to return colors as RGBA tuples in byte format (0 to 255)
+            - `ColorMap.FLOAT` or 'float' to return colors as RGBA tuples in float format (0.0 to 1.0)
+            - `ColorMap.QCOLOR` or 'qcolor' to return a list of QColors
+            
+        The default is byte format.
+        """
+        stops, color = self.getStops(mode=mode)
+        return color
+
+    def getStops(self, mode=BYTE):
+        """
+        Returns a tuple (stops, colors) containing a list of all stops (ranging 0.0 to 1.0)
+        and a list of the associated colors.
+        
+        The parameter `mode` can be one of
+            - `ColorMap.BYTE` or 'byte' to return colors as RGBA tuples in byte format (0 to 255)
+            - `ColorMap.FLOAT` or 'float' to return colors as RGBA tuples in float format (0.0 to 1.0)
+            - `ColorMap.QCOLOR` or 'qcolor' to return a list of QColors
+
+        The default is byte format.
         """
         if isinstance(mode, str):
             mode = self.enumMap[mode.lower()]
 
-        color = self.color
-        if mode in [self.BYTE, self.QCOLOR] and color.dtype.kind == 'f':
-            color = (color * 255).astype(np.ubyte)
-        elif mode == self.FLOAT and color.dtype.kind != 'f':
-            color = color.astype(float) / 255.
-
-        if mode == self.QCOLOR:
-            color = [QtGui.QColor(*x) for x in color]
-
-        return color
-
-    def getStops(self, mode):
-        ## Get fully-expanded set of RGBA stops in either float or byte mode.
         if mode not in self.stopsCache:
             color = self.color
             if mode == self.BYTE and color.dtype.kind == 'f':
                 color = (color*255).astype(np.ubyte)
             elif mode == self.FLOAT and color.dtype.kind != 'f':
                 color = color.astype(float) / 255.
+            elif mode == self.QCOLOR:
+                if color.dtype.kind == 'f':
+                    color = (color*255).astype(np.ubyte)
+                color = [QtGui.QColor(*x) for x in color]
             self.stopsCache[mode] = (self.pos, color)
         return self.stopsCache[mode]
 

tests/test_colormap.py

@@ -0,0 +1,75 @@
+import pytest
+import pyqtgraph as pg
+from pyqtgraph.Qt import QtGui
+
+pos = [0.0, 0.5, 1.0]
+qcols = [
+    QtGui.QColor('#FF0000'),
+    QtGui.QColor('#00FF00'),
+    QtGui.QColor('#0000FF')
+]
+float_tuples = [
+    (1.0, 0.0, 0.0, 1.0),
+    (0.0, 1.0, 0.0, 1.0),
+    (0.0, 0.0, 1.0, 1.0)
+]
+int_tuples = [
+    (255,  0,  0,255),
+    (  0,255,  0,255),
+    (  0,  0,255,255)
+] 
+
+@pytest.mark.parametrize("color_list", (qcols, int_tuples))
+def test_ColorMap_getStops(color_list):
+    cm = pg.ColorMap(pos, color_list, name='test')
+    # default is byte format:
+    stops, colors = cm.getStops()
+    assert (stops == pos).all()
+    assert (colors == int_tuples).all()
+
+    # manual byte format:
+    stops, colors = cm.getStops(pg.ColorMap.BYTE)
+    assert (stops == pos).all()
+    assert (colors == int_tuples).all()
+
+    stops, colors = cm.getStops('bYTe')
+    assert (stops == pos).all()
+    assert (colors == int_tuples).all()
+
+    # manual float format:
+    stops, colors = cm.getStops(pg.ColorMap.FLOAT)
+    assert (stops == pos).all()
+    assert (colors == float_tuples).all()
+
+    stops, colors = cm.getStops('floaT')
+    assert (stops == pos).all()
+    assert (colors == float_tuples).all()
+
+    # manual QColor format:
+    stops, colors = cm.getStops(pg.ColorMap.QCOLOR)
+    assert (stops == pos).all()
+    for actual, good in zip(colors, qcols):
+        assert actual.getRgbF() == good.getRgbF()
+
+    stops, colors = cm.getStops('qColor')
+    assert (stops == pos).all()
+    for actual, good in zip(colors, qcols):
+        assert actual.getRgbF() == good.getRgbF()
+
+
+@pytest.mark.parametrize("color_list", (qcols, int_tuples))
+def test_ColorMap_getColors(color_list):
+    cm = pg.ColorMap(pos, color_list, name='from QColors')
+
+    colors = cm.getColors()
+    assert (colors == int_tuples).all()
+
+    colors = cm.getColors('byte')
+    assert (colors == int_tuples).all()
+
+    colors = cm.getColors('float')
+    assert (colors == float_tuples).all()
+
+    colors = cm.getColors('qcolor')
+    for actual, good in zip(colors, qcols):
+        assert actual.getRgbF() == good.getRgbF()