diff --git a/branca/colormap.py b/branca/colormap.py index 2e76ba5..0418ec0 100644 --- a/branca/colormap.py +++ b/branca/colormap.py @@ -9,7 +9,7 @@ import json import math import os -from typing import Dict, List, Optional, Sequence, Tuple, Union +from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union from jinja2 import Template @@ -45,19 +45,20 @@ def _parse_hex(color_code: str) -> TypeRGBAFloats: ) -def _color_int_to_float(x: int) -> float: - """Convert an integer between 0 and 255 to a float between 0. and 1.0""" +def _color_int_to_float(x: Union[int, float]) -> float: + """Convert a byte between 0 and 255 to a normalized float between 0. and 1.0""" return x / 255.0 def _color_float_to_int(x: float) -> int: - """Convert a float between 0. and 1.0 to an integer between 0 and 255""" + """Convert a float between 0. and 1.0 to a byte integer between 0 and 255""" return int(x * 255.9999) def _parse_color(x: Union[tuple, list, str]) -> TypeRGBAFloats: + """Convert an unknown color value to an RGBA tuple of floats between 0 and 1.""" if isinstance(x, (tuple, list)): - return tuple(tuple(x) + (1.0,))[:4] # type: ignore + return _parse_color_as_numerical_sequence(x) elif isinstance(x, str) and _is_hex(x): return _parse_hex(x) elif isinstance(x, str): @@ -69,6 +70,32 @@ def _parse_color(x: Union[tuple, list, str]) -> TypeRGBAFloats: raise ValueError(f"Unrecognized color code {x!r}") +def _parse_color_as_numerical_sequence(x: Union[tuple, list]) -> TypeRGBAFloats: + """Convert a color as a sequence of numbers to an RGBA tuple of floats between 0 and 1.""" + if not all(isinstance(value, (int, float)) for value in x): + raise TypeError("Components in color sequence should all be int or float.") + if not 3 <= len(x) <= 4: + raise ValueError(f"Color sequence should have 3 or 4 components, not {len(x)}.") + if min(x) < 0 or max(x) > 255: + raise ValueError("Color components should be between 0.0 and 1.0 or 0 and 255.") + + if all(isinstance(value, int) for value in x): + # assume integers are a sequence of bytes that have to be normalized + conversion_function: Callable = _color_int_to_float + elif 1 < max(x) <= 255: + # values between 1 and 255 are bytes no matter the type and should be normalized + conversion_function = _color_int_to_float + else: + # else assume it's already normalized + conversion_function = float + + color: List[float] = [conversion_function(value) for value in x] + if len(color) == 3: + color.append(1.0) # add alpha channel + + return color[0], color[1], color[2], color[3] + + def _base(x: float) -> float: if x > 0: base = pow(10, math.floor(math.log10(x))) diff --git a/tests/test_colormap_parse.py b/tests/test_colormap_parse.py new file mode 100644 index 0000000..34e9559 --- /dev/null +++ b/tests/test_colormap_parse.py @@ -0,0 +1,145 @@ +import pytest + +from branca.colormap import ( + _color_float_to_int, + _color_int_to_float, + _is_hex, + _parse_color, + _parse_color_as_numerical_sequence, + _parse_hex, +) + + +@pytest.mark.parametrize( + "input_data, expected", + [ + ((255, 0, 0), (1.0, 0.0, 0.0, 1.0)), + ((255, 0, 0, 127), (1.0, 0.0, 0.0, 0.4980392156862745)), + ("#FF0000", (1.0, 0.0, 0.0, 1.0)), + ("red", (1.0, 0.0, 0.0, 1.0)), + ((0.5, 0.5, 0.5), (0.5, 0.5, 0.5, 1.0)), + ((0.25, 0.5, 0.75), (0.25, 0.5, 0.75, 1.0)), + ((0.1, 0.2, 0.3, 0.4), (0.1, 0.2, 0.3, 0.4)), + ("#0000FF", (0.0, 0.0, 1.0, 1.0)), + ("#00FF00", (0.0, 1.0, 0.0, 1.0)), + ("#FFFFFF", (1.0, 1.0, 1.0, 1.0)), + ("#000000", (0.0, 0.0, 0.0, 1.0)), + ("#808080", (0.5019607843137255, 0.5019607843137255, 0.5019607843137255, 1.0)), + ( + "#1A2B3C", + (0.10196078431372549, 0.16862745098039217, 0.23529411764705882, 1.0), + ), + ("green", (0.0, 0.5019607843137255, 0.0, 1.0)), + ], +) +def test_parse_color(input_data, expected): + assert _parse_color(input_data) == expected + + +@pytest.mark.parametrize( + "input_data, expected", + [ + # these are byte values as ints and should be normalized and converted + ((0, 0, 0), (0.0, 0.0, 0.0, 1.0)), + ((255, 255, 255), (1.0, 1.0, 1.0, 1.0)), + ((255, 0, 0), (1.0, 0.0, 0.0, 1.0)), + # a special case: ints that are 0 or 1 should be considered bytes + ((0, 0, 1), (0.0, 0.0, 1 / 255, 1.0)), + ((0, 0, 0, 1), (0.0, 0.0, 0.0, 1 / 255)), + # these already are normalized floats + ((0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0)), + ((0.0, 0.0, 1.0), (0.0, 0.0, 1.0, 1.0)), + ((0.5, 0.5, 0.5), (0.5, 0.5, 0.5, 1.0)), + ((0.1, 0.2, 0.3, 0.4), (0.1, 0.2, 0.3, 0.4)), + ((0.0, 1.0, 0.0, 0.5), (0.0, 1.0, 0.0, 0.5)), + # these are byte values as floats and should be normalized + ((0, 0, 0, 255.0), (0.0, 0.0, 0.0, 1.0)), + ((0, 0, 255.0, 0.0), (0.0, 0.0, 1.0, 0.0)), + # if floats and ints are mixed, assume they are intended as floats + ((0, 0, 1.0), (0.0, 0.0, 1.0, 1.0)), + # unless some of them are between 1 and 255 + ((0, 0, 1.0, 128), (0.0, 0.0, 1 / 255, 128 / 255)), + ], +) +def test_parse_color_as_numerical_sequence(input_data, expected): + assert _parse_color_as_numerical_sequence(input_data) == expected + + +@pytest.mark.parametrize( + "input_data, raises", + [ + # larger than 255 + ((256, 0, 0), ValueError), + # smaller than 0 + ((0, 0, -1), ValueError), + # sequence too long + ((0, 1, 2, 3, 4), ValueError), + # sequence too short + ((0, 1), ValueError), + # invalid type in sequence + ((0.5, 0.5, 0.5, "string"), TypeError), + ], +) +def test_parse_color_as_numerical_sequence_invalid(input_data, raises): + with pytest.raises(raises): + _parse_color_as_numerical_sequence(input_data) + + +@pytest.mark.parametrize( + "input_data, expected", + [ + ("#123456", True), + ("#abcdef", True), + ("#ABCDEF", True), + ("#1A2B3C", True), + ("#123", False), + ("123456", False), + ("#1234567", False), + ], +) +def test_is_hex(input_data, expected): + assert _is_hex(input_data) == expected + + +@pytest.mark.parametrize( + "input_data, expected", + [ + ("#000000", (0.0, 0.0, 0.0, 1.0)), + ("#FFFFFF", (1.0, 1.0, 1.0, 1.0)), + ("#FF0000", (1.0, 0.0, 0.0, 1.0)), + ("#00FF00", (0.0, 1.0, 0.0, 1.0)), + ("#0000FF", (0.0, 0.0, 1.0, 1.0)), + ("#808080", (0.5019607843137255, 0.5019607843137255, 0.5019607843137255, 1.0)), + ], +) +def test_parse_hex(input_data, expected): + assert _parse_hex(input_data) == expected + + +@pytest.mark.parametrize( + "input_data, expected", + [ + (0, 0.0), + (255, 1.0), + (128, 0.5019607843137255), + (64, 0.25098039215686274), + (192, 0.7529411764705882), + ], +) +def test_color_byte_to_normalized_float(input_data, expected): + assert _color_int_to_float(input_data) == expected + + +@pytest.mark.parametrize( + "input_data, expected", + [ + (0.0, 0), + (0.5, 127), + (1.0, 255), + (0.9999, 255), + (0.1, 25), + (0.75, 191), + ], +) +def test_color_normalized_float_to_byte_int(input_data, expected): + assert _color_float_to_int(input_data) == expected