Vehicle-cpp/hkpc/software/pandas/tests/plotting/frame/test_frame_legend.py

import numpy as np
import pytest

import pandas.util._test_decorators as td

from pandas import (
    DataFrame,
    date_range,
)
from pandas.tests.plotting.common import TestPlotBase
from pandas.util.version import Version


class TestFrameLegend(TestPlotBase):
    @pytest.mark.xfail(
        reason=(
            "Open bug in matplotlib "
            "https://github.com/matplotlib/matplotlib/issues/11357"
        )
    )
    def test_mixed_yerr(self):
        # https://github.com/pandas-dev/pandas/issues/39522
        import matplotlib as mpl
        from matplotlib.collections import LineCollection
        from matplotlib.lines import Line2D

        df = DataFrame([{"x": 1, "a": 1, "b": 1}, {"x": 2, "a": 2, "b": 3}])

        ax = df.plot("x", "a", c="orange", yerr=0.1, label="orange")
        df.plot("x", "b", c="blue", yerr=None, ax=ax, label="blue")

        legend = ax.get_legend()
        if Version(mpl.__version__) < Version("3.7"):
            result_handles = legend.legendHandles
        else:
            result_handles = legend.legend_handles

        assert isinstance(result_handles[0], LineCollection)
        assert isinstance(result_handles[1], Line2D)

    def test_legend_false(self):
        # https://github.com/pandas-dev/pandas/issues/40044
        import matplotlib as mpl

        df = DataFrame({"a": [1, 1], "b": [2, 3]})
        df2 = DataFrame({"d": [2.5, 2.5]})

        ax = df.plot(legend=True, color={"a": "blue", "b": "green"}, secondary_y="b")
        df2.plot(legend=True, color={"d": "red"}, ax=ax)
        legend = ax.get_legend()
        if Version(mpl.__version__) < Version("3.7"):
            handles = legend.legendHandles
        else:
            handles = legend.legend_handles
        result = [handle.get_color() for handle in handles]
        expected = ["blue", "green", "red"]
        assert result == expected

    @td.skip_if_no_scipy
    def test_df_legend_labels(self):
        kinds = ["line", "bar", "barh", "kde", "area", "hist"]
        df = DataFrame(np.random.rand(3, 3), columns=["a", "b", "c"])
        df2 = DataFrame(np.random.rand(3, 3), columns=["d", "e", "f"])
        df3 = DataFrame(np.random.rand(3, 3), columns=["g", "h", "i"])
        df4 = DataFrame(np.random.rand(3, 3), columns=["j", "k", "l"])

        for kind in kinds:
            ax = df.plot(kind=kind, legend=True)
            self._check_legend_labels(ax, labels=df.columns)

            ax = df2.plot(kind=kind, legend=False, ax=ax)
            self._check_legend_labels(ax, labels=df.columns)

            ax = df3.plot(kind=kind, legend=True, ax=ax)
            self._check_legend_labels(ax, labels=df.columns.union(df3.columns))

            ax = df4.plot(kind=kind, legend="reverse", ax=ax)
            expected = list(df.columns.union(df3.columns)) + list(reversed(df4.columns))
            self._check_legend_labels(ax, labels=expected)

        # Secondary Y
        ax = df.plot(legend=True, secondary_y="b")
        self._check_legend_labels(ax, labels=["a", "b (right)", "c"])
        ax = df2.plot(legend=False, ax=ax)
        self._check_legend_labels(ax, labels=["a", "b (right)", "c"])
        ax = df3.plot(kind="bar", legend=True, secondary_y="h", ax=ax)
        self._check_legend_labels(
            ax, labels=["a", "b (right)", "c", "g", "h (right)", "i"]
        )

        # Time Series
        ind = date_range("1/1/2014", periods=3)
        df = DataFrame(np.random.randn(3, 3), columns=["a", "b", "c"], index=ind)
        df2 = DataFrame(np.random.randn(3, 3), columns=["d", "e", "f"], index=ind)
        df3 = DataFrame(np.random.randn(3, 3), columns=["g", "h", "i"], index=ind)
        ax = df.plot(legend=True, secondary_y="b")
        self._check_legend_labels(ax, labels=["a", "b (right)", "c"])
        ax = df2.plot(legend=False, ax=ax)
        self._check_legend_labels(ax, labels=["a", "b (right)", "c"])
        ax = df3.plot(legend=True, ax=ax)
        self._check_legend_labels(ax, labels=["a", "b (right)", "c", "g", "h", "i"])

        # scatter
        ax = df.plot.scatter(x="a", y="b", label="data1")
        self._check_legend_labels(ax, labels=["data1"])
        ax = df2.plot.scatter(x="d", y="e", legend=False, label="data2", ax=ax)
        self._check_legend_labels(ax, labels=["data1"])
        ax = df3.plot.scatter(x="g", y="h", label="data3", ax=ax)
        self._check_legend_labels(ax, labels=["data1", "data3"])

        # ensure label args pass through and
        # index name does not mutate
        # column names don't mutate
        df5 = df.set_index("a")
        ax = df5.plot(y="b")
        self._check_legend_labels(ax, labels=["b"])
        ax = df5.plot(y="b", label="LABEL_b")
        self._check_legend_labels(ax, labels=["LABEL_b"])
        self._check_text_labels(ax.xaxis.get_label(), "a")
        ax = df5.plot(y="c", label="LABEL_c", ax=ax)
        self._check_legend_labels(ax, labels=["LABEL_b", "LABEL_c"])
        assert df5.columns.tolist() == ["b", "c"]

    def test_missing_marker_multi_plots_on_same_ax(self):
        # GH 18222
        df = DataFrame(data=[[1, 1, 1, 1], [2, 2, 4, 8]], columns=["x", "r", "g", "b"])
        fig, ax = self.plt.subplots(nrows=1, ncols=3)
        # Left plot
        df.plot(x="x", y="r", linewidth=0, marker="o", color="r", ax=ax[0])
        df.plot(x="x", y="g", linewidth=1, marker="x", color="g", ax=ax[0])
        df.plot(x="x", y="b", linewidth=1, marker="o", color="b", ax=ax[0])
        self._check_legend_labels(ax[0], labels=["r", "g", "b"])
        self._check_legend_marker(ax[0], expected_markers=["o", "x", "o"])
        # Center plot
        df.plot(x="x", y="b", linewidth=1, marker="o", color="b", ax=ax[1])
        df.plot(x="x", y="r", linewidth=0, marker="o", color="r", ax=ax[1])
        df.plot(x="x", y="g", linewidth=1, marker="x", color="g", ax=ax[1])
        self._check_legend_labels(ax[1], labels=["b", "r", "g"])
        self._check_legend_marker(ax[1], expected_markers=["o", "o", "x"])
        # Right plot
        df.plot(x="x", y="g", linewidth=1, marker="x", color="g", ax=ax[2])
        df.plot(x="x", y="b", linewidth=1, marker="o", color="b", ax=ax[2])
        df.plot(x="x", y="r", linewidth=0, marker="o", color="r", ax=ax[2])
        self._check_legend_labels(ax[2], labels=["g", "b", "r"])
        self._check_legend_marker(ax[2], expected_markers=["x", "o", "o"])

    def test_legend_name(self):
        multi = DataFrame(
            np.random.randn(4, 4),
            columns=[np.array(["a", "a", "b", "b"]), np.array(["x", "y", "x", "y"])],
        )
        multi.columns.names = ["group", "individual"]

        ax = multi.plot()
        leg_title = ax.legend_.get_title()
        self._check_text_labels(leg_title, "group,individual")

        df = DataFrame(np.random.randn(5, 5))
        ax = df.plot(legend=True, ax=ax)
        leg_title = ax.legend_.get_title()
        self._check_text_labels(leg_title, "group,individual")

        df.columns.name = "new"
        ax = df.plot(legend=False, ax=ax)
        leg_title = ax.legend_.get_title()
        self._check_text_labels(leg_title, "group,individual")

        ax = df.plot(legend=True, ax=ax)
        leg_title = ax.legend_.get_title()
        self._check_text_labels(leg_title, "new")

    @pytest.mark.parametrize(
        "kind",
        [
            "line",
            "bar",
            "barh",
            pytest.param("kde", marks=td.skip_if_no_scipy),
            "area",
            "hist",
        ],
    )
    def test_no_legend(self, kind):
        df = DataFrame(np.random.rand(3, 3), columns=["a", "b", "c"])
        ax = df.plot(kind=kind, legend=False)
        self._check_legend_labels(ax, visible=False)

    def test_missing_markers_legend(self):
        # 14958
        df = DataFrame(np.random.randn(8, 3), columns=["A", "B", "C"])
        ax = df.plot(y=["A"], marker="x", linestyle="solid")
        df.plot(y=["B"], marker="o", linestyle="dotted", ax=ax)
        df.plot(y=["C"], marker="<", linestyle="dotted", ax=ax)

        self._check_legend_labels(ax, labels=["A", "B", "C"])
        self._check_legend_marker(ax, expected_markers=["x", "o", "<"])

    def test_missing_markers_legend_using_style(self):
        # 14563
        df = DataFrame(
            {
                "A": [1, 2, 3, 4, 5, 6],
                "B": [2, 4, 1, 3, 2, 4],
                "C": [3, 3, 2, 6, 4, 2],
                "X": [1, 2, 3, 4, 5, 6],
            }
        )

        fig, ax = self.plt.subplots()
        for kind in "ABC":
            df.plot("X", kind, label=kind, ax=ax, style=".")

        self._check_legend_labels(ax, labels=["A", "B", "C"])
        self._check_legend_marker(ax, expected_markers=[".", ".", "."])