test_utilities.py

import json
import os
from pathlib import Path
from typing import List

import pytest

import branca.utilities as ut
from branca.colormap import LinearColormap

rootpath = Path(os.path.dirname(os.path.abspath(__file__))) / ".." / "branca"
color_brewer_minimum_n = 3
color_brewer_maximum_n = 253  # Why this limitation @ branca/utilities.py:108 ?


# Loads schemes and their meta-data
with open(rootpath / "_schemes.json") as f:
    schemes = json.loads(f.read())
with open(rootpath / "scheme_info.json") as f:
    scheme_info = json.loads(f.read())
with open(rootpath / "scheme_base_codes.json") as f:
    core_schemes = json.loads(f.read())["codes"]


def test_color_brewer_base():
    scheme = ut.color_brewer("YlGnBu", 9)
    assert scheme == [
        "#ffffd9",
        "#edf8b1",
        "#c7e9b4",
        "#7fcdbb",
        "#41b6c4",
        "#1d91c0",
        "#225ea8",
        "#253494",
        "#081d58",
    ]


def test_color_brewer_reverse():
    scheme = ut.color_brewer("YlGnBu")
    scheme_r = ut.color_brewer("YlGnBu_r")
    assert scheme[::-1] == scheme_r


def test_color_brewer_extendability():
    """
    The non-qualitative schemes should be extendable.

    :see https://github.com/python-visualization/branca/issues/104
    :see https://github.com/python-visualization/branca/issues/114

    Moreover, the following error was not reported via issues:
    * TypeError in the linear_gradient function when trying to extend any scheme.
    Indeed, in color_brewer, the key searched in the scheme database was not found,
    thus, it was passing `None` instead of a real scheme vector to linear_gradient.
    """
    for sname in core_schemes:
        for n in range(color_brewer_minimum_n, color_brewer_maximum_n + 1):
            try:
                scheme = ut.color_brewer(sname, n=n)
            except Exception as e:
                if scheme_info[sname] == "Qualitative" and isinstance(e, ValueError):
                    continue
                raise

            assert len(scheme) == n

            # When we try to extend a scheme,
            # the reverse is not always the exact reverse vector of the original one.
            # Thus, we do not test this property!
            _ = ut.color_brewer(sname + "_r", n=n)


def test_color_avoid_unexpected_error():
    """
    We had unexpected errors by providing some scheme name with unexpected value of `n`.
    This function tests them.

    Identified errors which was not reported via issues:
    * The scheme 'viridis' was not in the base_codes JSON;
    * Multiple scheme hadn't any metadata in scheme_info JSON;
    * When a `n` value provided to `color_scheme` was a float,
    it tried to select an unknown scheme without raising the right Exception type.
    """

    # Verify that every scheme has is present in base codes
    scheme_names = set()
    for sname in schemes.keys():
        scheme_names.add(sname.split("_")[0])
    assert scheme_names == set(core_schemes)

    # Verify that every scheme has a metadata
    assert scheme_names == set(scheme_info.keys())

    # Verify that every scheme can be generated in color_brewer using exotic value of `n`.
    # Note that big but valid values are generated by test_color_brewer_extendability.
    for sname in scheme_names:
        for n in (
            [-10]
            + list(range(-1, color_brewer_minimum_n))
            + list(range(color_brewer_maximum_n + 1, color_brewer_maximum_n + 10))
        ):
            with pytest.raises(ValueError):
                ut.color_brewer(sname, n)
        for n in [str(color_brewer_minimum_n), float(color_brewer_minimum_n), "abc"]:
            with pytest.raises(TypeError):
                ut.color_brewer(sname, n)


@pytest.mark.parametrize(
    "value,result",
    [
        (1, (1.0, "px")),
        ("1 px", (1.0, "px")),
        ("80 % ", (80.0, "%")),
        ("100%   ", (100.0, "%")),
        ("3 vw", (3.0, "vw")),
        ("3.14 rem", (3.14, "rem")),
        ((1, "px"), (1.0, "px")),
        ((80.0, "%"), (80.0, "%")),
    ],
)
def test_parse_size(value, result):
    assert ut._parse_size(value) == result


@pytest.mark.parametrize(
    "value",
    [
        "what?",
        "1.21 jigawatts",
        ut._parse_size,
        (1.21, 4.9),
    ],
)
def test_parse_size_exceptions(value):
    with pytest.raises((ValueError, TypeError)):
        ut._parse_size(value)


def test_write_png_mono():
    mono_image = [
        [0.24309289, 0.75997446, 0.02971671, 0.52830537],
        [0.62339252, 0.65369358, 0.41545387, 0.03307279],
    ]

    mono_png = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x04\x00\x00\x00\x02\x08\x06\x00\x00\x00\x7f\xa8}c\x00\x00\x00)IDATx\xdac\x08\x0c\x0c\xfc\x0f\x02\x9c\x9c\x9c\xff7n\xdc\xf8\x9f\xe1\xe2\xc5\x8b\xffo\xdf\xbe\xfd\xbf\xbb\xbb\xfb?77\xf7\x7f\x00f\x87\x14\xdd\x0c\r;\xc0\x00\x00\x00\x00IEND\xaeB`\x82"  # noqa E501
    assert ut.write_png(mono_image) == mono_png

    colormap = LinearColormap(colors=["red", "yellow", "green"])
    color_png = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x04\x00\x00\x00\x02\x08\x06\x00\x00\x00\x7f\xa8}c\x00\x00\x00)IDATx\xdac\xf8_\xcf\xf0\xbf\xea\x10\xc3\xff\xff\xfc\x0c\xff?\xfcg\xf8\xcfp\xe0\x19\xc3\xff\r\xf7\x80\x02\xb7\x80X\x90\xe1?\x00N\xca\x13\xcd\xfb\xad\r\xb8\x00\x00\x00\x00IEND\xaeB`\x82"  # noqa E501
    assert ut.write_png(mono_image, colormap=colormap) == color_png


def test_write_png_rgb():
    image_rgb = [
        [
            [0.8952778565195247, 0.6196806506704735, 0.2696137085302287],
            [0.3940794236804127, 0.9432178293916365, 0.16500617914697335],
            [0.5566755388192485, 0.10469673377265687, 0.27346260130585975],
            [0.2029951628162342, 0.5357152681832641, 0.13692921080346832],
        ],
        [
            [0.5186482474007286, 0.8625240370164696, 0.6965561989987038],
            [0.04425586727957387, 0.45448042432657076, 0.8552600511205423],
            [0.696453974598333, 0.7508742900711168, 0.9646572952994652],
            [0.7471809029502141, 0.3218907599994758, 0.789193070740859],
        ],
    ]
    png = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x04\x00\x00\x00\x02\x08\x06\x00\x00\x00\x7f\xa8}c\x00\x00\x00-IDATx\xda\x01"\x00\xdd\xff\x00\xff\xa7G\xffp\xff+\xff\x9e\x1cH\xff9\x90$\xff\x00\x93\xe9\xb8\xff\x0cz\xe2\xff\xc6\xca\xff\xff\xd4W\xd0\xffYw\x15\x95\xcf\xb9@D\x00\x00\x00\x00IEND\xaeB`\x82'  # noqa E501
    assert ut.write_png(image_rgb) == png


@pytest.mark.parametrize(
    "hex_list, n_colors, expected_output",
    [
        (["#000000", "#FFFFFF"], 2, ["#000000", "#ffffff"]),
        (["#000000", "#FFFFFF"], 4, ["#000000", "#545454", "#a9a9a9", "#ffffff"]),
        (["#FF0000", "#00FF00", "#0000FF"], 3, ["#ff0000", "#00ff00", "#0000ff"]),
        (
            ["#FF0000", "#00FF00", "#0000FF"],
            4,
            ["#ff0000", "#55a900", "#00aa54", "#0000ff"],
        ),
        (
            ["#000000", "#0000FF"],
            5,
            ["#000000", "#00003f", "#00007f", "#0000bf", "#0000ff"],
        ),
        (
            ["#FFFFFF", "#000000"],
            5,
            ["#ffffff", "#bfbfbf", "#7f7f7f", "#3f3f3f", "#000000"],
        ),
        (
            ["#FF0000", "#00FF00", "#0000FF"],
            5,
            ["#ff0000", "#7f7f00", "#00ff00", "#007f7f", "#0000ff"],
        ),
        (
            ["#FF0000", "#00FF00", "#0000FF"],
            7,
            [
                "#ff0000",
                "#aa5400",
                "#55a900",
                "#00ff00",
                "#00aa54",
                "#0055a9",
                "#0000ff",
            ],
        ),
        (
            ["#abcdef", "#010603", "#f7f9f3"],
            4,
            ["#abcdef", "#394851", "#525652", "#f7f9f3"],
        ),
        (
            ["#abcdef", "#010603", "#f7f9f3"],
            7,
            [
                "#abcdef",
                "#728aa0",
                "#394851",
                "#010603",
                "#525652",
                "#a4a7a2",
                "#f7f9f3",
            ],
        ),
        (
            ["#00abff", "#ff00ab", "#abff00", "#00abff"],
            4,
            ["#00abff", "#ff00ab", "#abff00", "#00abff"],
        ),
        (
            ["#00abff", "#ff00ab", "#abff00", "#00abff"],
            6,
            ["#00abff", "#9844cc", "#ee3288", "#bbcb22", "#66dd65", "#00abff"],
        ),
    ],
)
def test_linear_gradient(
    hex_list: List[str],
    n_colors: int,
    expected_output: List[str],
):
    result = ut.linear_gradient(hex_list, n_colors)
    assert result == expected_output