From f79d0dfa1490dea075800ac04b1b4701c597a73f Mon Sep 17 00:00:00 2001 From: Nils Nemitz Date: Tue, 22 Jun 2021 01:19:07 +0900 Subject: [PATCH] make getColors and getStops behavior more consistent; add tests --- pyqtgraph/colormap.py | 46 +++++++++++++++++--------- tests/test_colormap.py | 75 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 16 deletions(-) create mode 100644 tests/test_colormap.py diff --git a/pyqtgraph/colormap.py b/pyqtgraph/colormap.py index 5f58da67..a736a46e 100644 --- a/pyqtgraph/colormap.py +++ b/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,44 @@ 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 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-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-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] diff --git a/tests/test_colormap.py b/tests/test_colormap.py new file mode 100644 index 00000000..35e62d34 --- /dev/null +++ b/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()