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 (
_check_legend_labels,
_check_legend_marker,
_check_text_labels,
)
from pandas.util.version import Version
mpl = pytest.importorskip("matplotlib")
class TestFrameLegend:
@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
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
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
@pytest.mark.parametrize("kind", ["line", "bar", "barh", "kde", "area", "hist"])
def test_df_legend_labels(self, kind):
pytest.importorskip("scipy")
df = DataFrame(np.random.default_rng(2).random((3, 3)), columns=["a", "b", "c"])
df2 = DataFrame(
np.random.default_rng(2).random((3, 3)), columns=["d", "e", "f"]
)
df3 = DataFrame(
np.random.default_rng(2).random((3, 3)), columns=["g", "h", "i"]
)
df4 = DataFrame(
np.random.default_rng(2).random((3, 3)), columns=["j", "k", "l"]
)
ax = df.plot(kind=kind, legend=True)
_check_legend_labels(ax, labels=df.columns)
ax = df2.plot(kind=kind, legend=False, ax=ax)
_check_legend_labels(ax, labels=df.columns)
ax = df3.plot(kind=kind, legend=True, ax=ax)
_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))
_check_legend_labels(ax, labels=expected)
def test_df_legend_labels_secondary_y(self):
pytest.importorskip("scipy")
df = DataFrame(np.random.default_rng(2).random((3, 3)), columns=["a", "b", "c"])
df2 = DataFrame(
np.random.default_rng(2).random((3, 3)), columns=["d", "e", "f"]
)
df3 = DataFrame(
np.random.default_rng(2).random((3, 3)), columns=["g", "h", "i"]
)
# Secondary Y
ax = df.plot(legend=True, secondary_y="b")
_check_legend_labels(ax, labels=["a", "b (right)", "c"])
ax = df2.plot(legend=False, ax=ax)
_check_legend_labels(ax, labels=["a", "b (right)", "c"])
ax = df3.plot(kind="bar", legend=True, secondary_y="h", ax=ax)
_check_legend_labels(ax, labels=["a", "b (right)", "c", "g", "h (right)", "i"])
def test_df_legend_labels_time_series(self):
# Time Series
pytest.importorskip("scipy")
ind = date_range("1/1/2014", periods=3)
df = DataFrame(
np.random.default_rng(2).standard_normal((3, 3)),
columns=["a", "b", "c"],
index=ind,
)
df2 = DataFrame(
np.random.default_rng(2).standard_normal((3, 3)),
columns=["d", "e", "f"],
index=ind,
)
df3 = DataFrame(
np.random.default_rng(2).standard_normal((3, 3)),
columns=["g", "h", "i"],
index=ind,
)
ax = df.plot(legend=True, secondary_y="b")
_check_legend_labels(ax, labels=["a", "b (right)", "c"])
ax = df2.plot(legend=False, ax=ax)
_check_legend_labels(ax, labels=["a", "b (right)", "c"])
ax = df3.plot(legend=True, ax=ax)
_check_legend_labels(ax, labels=["a", "b (right)", "c", "g", "h", "i"])
def test_df_legend_labels_time_series_scatter(self):
# Time Series
pytest.importorskip("scipy")
ind = date_range("1/1/2014", periods=3)
df = DataFrame(
np.random.default_rng(2).standard_normal((3, 3)),
columns=["a", "b", "c"],
index=ind,
)
df2 = DataFrame(
np.random.default_rng(2).standard_normal((3, 3)),
columns=["d", "e", "f"],
index=ind,
)
df3 = DataFrame(
np.random.default_rng(2).standard_normal((3, 3)),
columns=["g", "h", "i"],
index=ind,
)
# scatter
ax = df.plot.scatter(x="a", y="b", label="data1")
_check_legend_labels(ax, labels=["data1"])
ax = df2.plot.scatter(x="d", y="e", legend=False, label="data2", ax=ax)
_check_legend_labels(ax, labels=["data1"])
ax = df3.plot.scatter(x="g", y="h", label="data3", ax=ax)
_check_legend_labels(ax, labels=["data1", "data3"])
def test_df_legend_labels_time_series_no_mutate(self):
pytest.importorskip("scipy")
ind = date_range("1/1/2014", periods=3)
df = DataFrame(
np.random.default_rng(2).standard_normal((3, 3)),
columns=["a", "b", "c"],
index=ind,
)
# 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")
_check_legend_labels(ax, labels=["b"])
ax = df5.plot(y="b", label="LABEL_b")
_check_legend_labels(ax, labels=["LABEL_b"])
_check_text_labels(ax.xaxis.get_label(), "a")
ax = df5.plot(y="c", label="LABEL_c", ax=ax)
_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"])
_, ax = mpl.pyplot.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])
_check_legend_labels(ax[0], labels=["r", "g", "b"])
_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])
_check_legend_labels(ax[1], labels=["b", "r", "g"])
_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])
_check_legend_labels(ax[2], labels=["g", "b", "r"])
_check_legend_marker(ax[2], expected_markers=["x", "o", "o"])
def test_legend_name(self):
multi = DataFrame(
np.random.default_rng(2).standard_normal((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()
_check_text_labels(leg_title, "group,individual")
df = DataFrame(np.random.default_rng(2).standard_normal((5, 5)))
ax = df.plot(legend=True, ax=ax)
leg_title = ax.legend_.get_title()
_check_text_labels(leg_title, "group,individual")
df.columns.name = "new"
ax = df.plot(legend=False, ax=ax)
leg_title = ax.legend_.get_title()
_check_text_labels(leg_title, "group,individual")
ax = df.plot(legend=True, ax=ax)
leg_title = ax.legend_.get_title()
_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.default_rng(2).random((3, 3)), columns=["a", "b", "c"])
ax = df.plot(kind=kind, legend=False)
_check_legend_labels(ax, visible=False)
def test_missing_markers_legend(self):
# 14958
df = DataFrame(
np.random.default_rng(2).standard_normal((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)
_check_legend_labels(ax, labels=["A", "B", "C"])
_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],
}
)
_, ax = mpl.pyplot.subplots()
for kind in "ABC":
df.plot("X", kind, label=kind, ax=ax, style=".")
_check_legend_labels(ax, labels=["A", "B", "C"])
_check_legend_marker(ax, expected_markers=[".", ".", "."])
|