ETPlot/tests/test_plotstyle.py

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import pytest
from matplotlib.colors import Colormap

import plotstyle as ps
from plotstyle import colors


def test_categorical_palette_matches_kit_hex():
    assert colors.CATEGORICAL == [
        "#009682",
        "#DF9B1B",
        "#4664AA",
        "#A78230",
        "#23A1E0",
        "#A3107C",
        "#8CB63C",
        "#A22223",
        "#FCE500",
    ]


def test_diverging_palette_matches_kit_hex():
    assert colors.DIVERGING == {"low": "#4664AA", "mid": "#f7f7f7", "high": "#A22223"}


def test_status_palette_matches_validated_hex():
    assert colors.STATUS == {
        "good": "#0ca30c",
        "warning": "#fab219",
        "serious": "#ec835a",
        "critical": "#d03b3b",
    }


def test_get_color_returns_categorical_slot():
    assert colors.get_color(0) == colors.CATEGORICAL[0]
    assert colors.get_color(len(colors.CATEGORICAL) - 1) == colors.CATEGORICAL[-1]


def test_get_color_out_of_range_raises():
    with pytest.raises(ValueError):
        colors.get_color(len(colors.CATEGORICAL))


def test_sequential_steps_are_kit_blue_variations():
    assert colors.SEQUENTIAL_STEPS[-1].lower() == "#4664aa"
    # Lightest step should be a lighter (higher-lightness) tint than the brand color.
    import colorsys

    def lightness(hex_color):
        r, g, b = (int(hex_color.lstrip("#")[i : i + 2], 16) / 255 for i in (0, 2, 4))
        return colorsys.rgb_to_hls(r, g, b)[1]

    assert lightness(colors.SEQUENTIAL_STEPS[0]) > lightness(colors.SEQUENTIAL_STEPS[-1])


def test_sequential_cmap_is_usable_colormap():
    cmap = colors.sequential_cmap()
    assert isinstance(cmap, Colormap)
    assert cmap(0.0) != cmap(1.0)


def test_diverging_cmap_is_usable_colormap():
    cmap = colors.diverging_cmap()
    assert isinstance(cmap, Colormap)
    assert cmap(0.0) != cmap(1.0)


def test_use_sets_expected_rcparams():
    ps.use()
    try:
        assert plt.rcParams["axes.facecolor"] == "#ffffff"
        assert plt.rcParams["text.usetex"] is True
        assert "lmodern" in plt.rcParams["text.latex.preamble"]
        assert "sfmath" in plt.rcParams["text.latex.preamble"]
        assert "fontenc" in plt.rcParams["text.latex.preamble"]
        assert r"\sfdefault" in plt.rcParams["text.latex.preamble"]
        cycle_colors = [entry["color"] for entry in plt.rcParams["axes.prop_cycle"]]
        assert cycle_colors == colors.CATEGORICAL

        assert plt.rcParams["axes.spines.top"] is False
        assert plt.rcParams["axes.spines.right"] is False
        assert plt.rcParams["axes.spines.left"] is False
        assert plt.rcParams["axes.spines.bottom"] is True
        assert plt.rcParams["xtick.bottom"] is True
        assert plt.rcParams["ytick.left"] is True
        assert plt.rcParams["axes.edgecolor"] == "#000000"
        assert plt.rcParams["axes.linewidth"] > 0.8

        assert plt.rcParams["axes.grid"] is True
        assert plt.rcParams["axes.grid.axis"] == "y"
        assert plt.rcParams["axes.titlelocation"] == "left"

        assert plt.rcParams["xtick.minor.visible"] is True
        assert plt.rcParams["ytick.minor.visible"] is False
        assert plt.rcParams["xtick.major.size"] > plt.rcParams["xtick.minor.size"]
    finally:
        ps.reset()


def test_use_default_does_not_cycle_linestyles():
    ps.use()
    try:
        cycle_keys = plt.rcParams["axes.prop_cycle"].keys
        assert cycle_keys == {"color"}
    finally:
        ps.reset()


def test_use_cycle_linestyles_opt_in():
    ps.use(cycle_linestyles=True)
    try:
        cycle_keys = plt.rcParams["axes.prop_cycle"].keys
        assert cycle_keys == {"color", "linestyle"}
        linestyles = {entry["linestyle"] for entry in plt.rcParams["axes.prop_cycle"]}
        assert len(linestyles) > 1
    finally:
        ps.reset()


def test_new_figure_styles_major_and_minor_tick_labels_differently():
    fig, ax = ps.new_figure("square")
    try:
        assert ax.xaxis.get_tick_params(which="major")["labelcolor"] == colors.INK["primary"]
        assert ax.xaxis.get_tick_params(which="minor")["labelcolor"] == colors.INK["muted"]
        assert ax.yaxis.get_tick_params(which="major")["labelcolor"] == colors.INK["primary"]
        assert ax.yaxis.get_tick_params(which="minor")["labelcolor"] == colors.INK["muted"]
    finally:
        plt.close(fig)


def test_new_figure_preset_sizes():
    fig, ax = ps.new_figure("thesis-single")
    try:
        assert tuple(fig.get_size_inches()) == ps.FIGSIZES["thesis-single"]
    finally:
        plt.close(fig)


def test_new_figure_explicit_tuple():
    fig, ax = ps.new_figure((3.0, 2.0))
    try:
        assert tuple(fig.get_size_inches()) == (3.0, 2.0)
    finally:
        plt.close(fig)


def test_new_figure_unknown_preset_raises():
    with pytest.raises(ValueError):
        ps.new_figure("not-a-real-preset")


def test_new_figure_title_sets_left_aligned_figure_suptitle():
    fig, ax = ps.new_figure("square", title="My Title")
    try:
        assert fig._suptitle is not None
        assert "My Title" in fig._suptitle.get_text()
        assert fig._suptitle.get_ha() == "left"
        assert ax.get_title() == ""
    finally:
        plt.close(fig)


def test_new_figure_params_adds_subtitle_line():
    fig, ax = ps.new_figure("square", title="My Title", params={"N": 100, "seed": 42})
    try:
        text = fig._suptitle.get_text()
        assert "My Title" in text
        assert "N: 100" in text
        assert "seed: 42" in text
        assert "|" in text
    finally:
        plt.close(fig)


def test_new_figure_without_title_or_params_has_no_suptitle():
    fig, ax = ps.new_figure("square")
    try:
        assert fig._suptitle is None
    finally:
        plt.close(fig)


def test_colorbar_matches_axes_height():
    fig, ax = ps.new_figure("square")
    try:
        im = ax.imshow([[0, 1], [2, 3]])
        cb = ps.colorbar(im, ax)
        fig.canvas.draw()
        ax_bbox = ax.get_position()
        cax_bbox = cb.ax.get_position()
        assert ax_bbox.y0 == pytest.approx(cax_bbox.y0, abs=1e-6)
        assert ax_bbox.y1 == pytest.approx(cax_bbox.y1, abs=1e-6)
    finally:
        plt.close(fig)


def test_colorbar_has_no_outline():
    fig, ax = ps.new_figure("square")
    try:
        im = ax.imshow([[0, 1], [2, 3]])
        cb = ps.colorbar(im, ax)
        assert cb.outline.get_visible() is False
    finally:
        plt.close(fig)


def test_savefig_writes_requested_formats(tmp_path):
    fig, ax = ps.new_figure("square")
    ax.plot([0, 1], [0, 1])
    try:
        out_dir = tmp_path / "nested" / "plots"
        written = ps.savefig(fig, out_dir / "my_plot", formats=("pdf", "png"))
    finally:
        plt.close(fig)

    assert [p.name for p in written] == ["my_plot.pdf", "my_plot.png"]
    for p in written:
        assert p.exists()
        assert p.stat().st_size > 0


def test_style_legend_places_legend_on_figure_outside_axes():
    fig, ax = ps.new_figure("square")
    try:
        ax.plot([0, 1], [0, 1], label="series")
        legend = ps.style_legend(ax, title="Series")
        assert legend in fig.legends
        assert ax.get_legend() is None
        assert legend.get_title().get_text() == "Series"
    finally:
        plt.close(fig)


def test_style_legend_without_title_warns():
    fig, ax = ps.new_figure("square")
    try:
        ax.plot([0, 1], [0, 1], label="series")
        with pytest.warns(UserWarning):
            ps.style_legend(ax)
    finally:
        plt.close(fig)


def test_panel_label_does_not_raise():
    fig, ax = ps.new_figure("square")
    try:
        ax.plot([0, 1], [0, 1], label="series")
        ps.panel_label(ax, "a")
    finally:
        plt.close(fig)


def test_panel_label_unknown_loc_raises():
    fig, ax = ps.new_figure("square")
    try:
        with pytest.raises(ValueError):
            ps.panel_label(ax, "a", loc="middle")
    finally:
        plt.close(fig)


def test_panel_label_defaults_to_lower_right():
    fig, ax = ps.new_figure("square")
    try:
        text = ps.panel_label(ax, "a")
        assert text.get_ha() == "right"
        assert text.get_va() == "bottom"
        x, y = text.get_position()
        assert x > 0.5
        assert y < 0.5
    finally:
        plt.close(fig)


def test_panel_label_color_matches_xlabel():
    fig, ax = ps.new_figure("square")
    try:
        ax.set_xlabel("Time (s)")
        text = ps.panel_label(ax, "a")
        assert text.get_color() == ax.xaxis.label.get_color()
    finally:
        plt.close(fig)


def test_panel_label_has_rounded_semi_transparent_box_by_default():
    fig, ax = ps.new_figure("square")
    try:
        text = ps.panel_label(ax, "a")
        patch = text.get_bbox_patch()
        assert patch is not None
        assert "round" in patch.get_boxstyle().__class__.__name__.lower()
        assert patch.get_facecolor()[3] < 1.0
        assert patch.get_edgecolor()[3] == 1.0
    finally:
        plt.close(fig)


def test_panel_label_box_can_be_disabled():
    fig, ax = ps.new_figure("square")
    try:
        text = ps.panel_label(ax, "a", box=False)
        assert text.get_bbox_patch() is None
    finally:
        plt.close(fig)