215 lines
8.5 KiB
Python
215 lines
8.5 KiB
Python
|
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=[".", ".", "."])
|