| """ |
| Module consolidating common testing functions for checking plotting. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING |
|
|
| import numpy as np |
|
|
| from pandas.core.dtypes.api import is_list_like |
|
|
| import pandas as pd |
| from pandas import Series |
| import pandas._testing as tm |
|
|
| if TYPE_CHECKING: |
| from collections.abc import Sequence |
|
|
| from matplotlib.axes import Axes |
|
|
|
|
| def _check_legend_labels(axes, labels=None, visible=True): |
| """ |
| Check each axes has expected legend labels |
| |
| Parameters |
| ---------- |
| axes : matplotlib Axes object, or its list-like |
| labels : list-like |
| expected legend labels |
| visible : bool |
| expected legend visibility. labels are checked only when visible is |
| True |
| """ |
| if visible and (labels is None): |
| raise ValueError("labels must be specified when visible is True") |
| axes = _flatten_visible(axes) |
| for ax in axes: |
| if visible: |
| assert ax.get_legend() is not None |
| _check_text_labels(ax.get_legend().get_texts(), labels) |
| else: |
| assert ax.get_legend() is None |
|
|
|
|
| def _check_legend_marker(ax, expected_markers=None, visible=True): |
| """ |
| Check ax has expected legend markers |
| |
| Parameters |
| ---------- |
| ax : matplotlib Axes object |
| expected_markers : list-like |
| expected legend markers |
| visible : bool |
| expected legend visibility. labels are checked only when visible is |
| True |
| """ |
| if visible and (expected_markers is None): |
| raise ValueError("Markers must be specified when visible is True") |
| if visible: |
| handles, _ = ax.get_legend_handles_labels() |
| markers = [handle.get_marker() for handle in handles] |
| assert markers == expected_markers |
| else: |
| assert ax.get_legend() is None |
|
|
|
|
| def _check_data(xp, rs): |
| """ |
| Check each axes has identical lines |
| |
| Parameters |
| ---------- |
| xp : matplotlib Axes object |
| rs : matplotlib Axes object |
| """ |
| import matplotlib.pyplot as plt |
|
|
| xp_lines = xp.get_lines() |
| rs_lines = rs.get_lines() |
|
|
| assert len(xp_lines) == len(rs_lines) |
| for xpl, rsl in zip(xp_lines, rs_lines): |
| xpdata = xpl.get_xydata() |
| rsdata = rsl.get_xydata() |
| tm.assert_almost_equal(xpdata, rsdata) |
|
|
| plt.close("all") |
|
|
|
|
| def _check_visible(collections, visible=True): |
| """ |
| Check each artist is visible or not |
| |
| Parameters |
| ---------- |
| collections : matplotlib Artist or its list-like |
| target Artist or its list or collection |
| visible : bool |
| expected visibility |
| """ |
| from matplotlib.collections import Collection |
|
|
| if not isinstance(collections, Collection) and not is_list_like(collections): |
| collections = [collections] |
|
|
| for patch in collections: |
| assert patch.get_visible() == visible |
|
|
|
|
| def _check_patches_all_filled(axes: Axes | Sequence[Axes], filled: bool = True) -> None: |
| """ |
| Check for each artist whether it is filled or not |
| |
| Parameters |
| ---------- |
| axes : matplotlib Axes object, or its list-like |
| filled : bool |
| expected filling |
| """ |
|
|
| axes = _flatten_visible(axes) |
| for ax in axes: |
| for patch in ax.patches: |
| assert patch.fill == filled |
|
|
|
|
| def _get_colors_mapped(series, colors): |
| unique = series.unique() |
| |
| |
| mapped = dict(zip(unique, colors)) |
| return [mapped[v] for v in series.values] |
|
|
|
|
| def _check_colors(collections, linecolors=None, facecolors=None, mapping=None): |
| """ |
| Check each artist has expected line colors and face colors |
| |
| Parameters |
| ---------- |
| collections : list-like |
| list or collection of target artist |
| linecolors : list-like which has the same length as collections |
| list of expected line colors |
| facecolors : list-like which has the same length as collections |
| list of expected face colors |
| mapping : Series |
| Series used for color grouping key |
| used for andrew_curves, parallel_coordinates, radviz test |
| """ |
| from matplotlib import colors |
| from matplotlib.collections import ( |
| Collection, |
| LineCollection, |
| PolyCollection, |
| ) |
| from matplotlib.lines import Line2D |
|
|
| conv = colors.ColorConverter |
| if linecolors is not None: |
| if mapping is not None: |
| linecolors = _get_colors_mapped(mapping, linecolors) |
| linecolors = linecolors[: len(collections)] |
|
|
| assert len(collections) == len(linecolors) |
| for patch, color in zip(collections, linecolors): |
| if isinstance(patch, Line2D): |
| result = patch.get_color() |
| |
| result = conv.to_rgba(result) |
| elif isinstance(patch, (PolyCollection, LineCollection)): |
| result = tuple(patch.get_edgecolor()[0]) |
| else: |
| result = patch.get_edgecolor() |
|
|
| expected = conv.to_rgba(color) |
| assert result == expected |
|
|
| if facecolors is not None: |
| if mapping is not None: |
| facecolors = _get_colors_mapped(mapping, facecolors) |
| facecolors = facecolors[: len(collections)] |
|
|
| assert len(collections) == len(facecolors) |
| for patch, color in zip(collections, facecolors): |
| if isinstance(patch, Collection): |
| |
| result = patch.get_facecolor()[0] |
| else: |
| result = patch.get_facecolor() |
|
|
| if isinstance(result, np.ndarray): |
| result = tuple(result) |
|
|
| expected = conv.to_rgba(color) |
| assert result == expected |
|
|
|
|
| def _check_text_labels(texts, expected): |
| """ |
| Check each text has expected labels |
| |
| Parameters |
| ---------- |
| texts : matplotlib Text object, or its list-like |
| target text, or its list |
| expected : str or list-like which has the same length as texts |
| expected text label, or its list |
| """ |
| if not is_list_like(texts): |
| assert texts.get_text() == expected |
| else: |
| labels = [t.get_text() for t in texts] |
| assert len(labels) == len(expected) |
| for label, e in zip(labels, expected): |
| assert label == e |
|
|
|
|
| def _check_ticks_props(axes, xlabelsize=None, xrot=None, ylabelsize=None, yrot=None): |
| """ |
| Check each axes has expected tick properties |
| |
| Parameters |
| ---------- |
| axes : matplotlib Axes object, or its list-like |
| xlabelsize : number |
| expected xticks font size |
| xrot : number |
| expected xticks rotation |
| ylabelsize : number |
| expected yticks font size |
| yrot : number |
| expected yticks rotation |
| """ |
| from matplotlib.ticker import NullFormatter |
|
|
| axes = _flatten_visible(axes) |
| for ax in axes: |
| if xlabelsize is not None or xrot is not None: |
| if isinstance(ax.xaxis.get_minor_formatter(), NullFormatter): |
| |
| |
| labels = ax.get_xticklabels() |
| else: |
| labels = ax.get_xticklabels() + ax.get_xticklabels(minor=True) |
|
|
| for label in labels: |
| if xlabelsize is not None: |
| tm.assert_almost_equal(label.get_fontsize(), xlabelsize) |
| if xrot is not None: |
| tm.assert_almost_equal(label.get_rotation(), xrot) |
|
|
| if ylabelsize is not None or yrot is not None: |
| if isinstance(ax.yaxis.get_minor_formatter(), NullFormatter): |
| labels = ax.get_yticklabels() |
| else: |
| labels = ax.get_yticklabels() + ax.get_yticklabels(minor=True) |
|
|
| for label in labels: |
| if ylabelsize is not None: |
| tm.assert_almost_equal(label.get_fontsize(), ylabelsize) |
| if yrot is not None: |
| tm.assert_almost_equal(label.get_rotation(), yrot) |
|
|
|
|
| def _check_ax_scales(axes, xaxis="linear", yaxis="linear"): |
| """ |
| Check each axes has expected scales |
| |
| Parameters |
| ---------- |
| axes : matplotlib Axes object, or its list-like |
| xaxis : {'linear', 'log'} |
| expected xaxis scale |
| yaxis : {'linear', 'log'} |
| expected yaxis scale |
| """ |
| axes = _flatten_visible(axes) |
| for ax in axes: |
| assert ax.xaxis.get_scale() == xaxis |
| assert ax.yaxis.get_scale() == yaxis |
|
|
|
|
| def _check_axes_shape(axes, axes_num=None, layout=None, figsize=None): |
| """ |
| Check expected number of axes is drawn in expected layout |
| |
| Parameters |
| ---------- |
| axes : matplotlib Axes object, or its list-like |
| axes_num : number |
| expected number of axes. Unnecessary axes should be set to |
| invisible. |
| layout : tuple |
| expected layout, (expected number of rows , columns) |
| figsize : tuple |
| expected figsize. default is matplotlib default |
| """ |
| from pandas.plotting._matplotlib.tools import flatten_axes |
|
|
| if figsize is None: |
| figsize = (6.4, 4.8) |
| visible_axes = _flatten_visible(axes) |
|
|
| if axes_num is not None: |
| assert len(visible_axes) == axes_num |
| for ax in visible_axes: |
| |
| assert len(ax.get_children()) > 0 |
|
|
| if layout is not None: |
| x_set = set() |
| y_set = set() |
| for ax in flatten_axes(axes): |
| |
| points = ax.get_position().get_points() |
| x_set.add(points[0][0]) |
| y_set.add(points[0][1]) |
| result = (len(y_set), len(x_set)) |
| assert result == layout |
|
|
| tm.assert_numpy_array_equal( |
| visible_axes[0].figure.get_size_inches(), |
| np.array(figsize, dtype=np.float64), |
| ) |
|
|
|
|
| def _flatten_visible(axes: Axes | Sequence[Axes]) -> Sequence[Axes]: |
| """ |
| Flatten axes, and filter only visible |
| |
| Parameters |
| ---------- |
| axes : matplotlib Axes object, or its list-like |
| |
| """ |
| from pandas.plotting._matplotlib.tools import flatten_axes |
|
|
| axes_ndarray = flatten_axes(axes) |
| axes = [ax for ax in axes_ndarray if ax.get_visible()] |
| return axes |
|
|
|
|
| def _check_has_errorbars(axes, xerr=0, yerr=0): |
| """ |
| Check axes has expected number of errorbars |
| |
| Parameters |
| ---------- |
| axes : matplotlib Axes object, or its list-like |
| xerr : number |
| expected number of x errorbar |
| yerr : number |
| expected number of y errorbar |
| """ |
| axes = _flatten_visible(axes) |
| for ax in axes: |
| containers = ax.containers |
| xerr_count = 0 |
| yerr_count = 0 |
| for c in containers: |
| has_xerr = getattr(c, "has_xerr", False) |
| has_yerr = getattr(c, "has_yerr", False) |
| if has_xerr: |
| xerr_count += 1 |
| if has_yerr: |
| yerr_count += 1 |
| assert xerr == xerr_count |
| assert yerr == yerr_count |
|
|
|
|
| def _check_box_return_type( |
| returned, return_type, expected_keys=None, check_ax_title=True |
| ): |
| """ |
| Check box returned type is correct |
| |
| Parameters |
| ---------- |
| returned : object to be tested, returned from boxplot |
| return_type : str |
| return_type passed to boxplot |
| expected_keys : list-like, optional |
| group labels in subplot case. If not passed, |
| the function checks assuming boxplot uses single ax |
| check_ax_title : bool |
| Whether to check the ax.title is the same as expected_key |
| Intended to be checked by calling from ``boxplot``. |
| Normal ``plot`` doesn't attach ``ax.title``, it must be disabled. |
| """ |
| from matplotlib.axes import Axes |
|
|
| types = {"dict": dict, "axes": Axes, "both": tuple} |
| if expected_keys is None: |
| |
| if return_type is None: |
| return_type = "dict" |
|
|
| assert isinstance(returned, types[return_type]) |
| if return_type == "both": |
| assert isinstance(returned.ax, Axes) |
| assert isinstance(returned.lines, dict) |
| else: |
| |
| if return_type is None: |
| for r in _flatten_visible(returned): |
| assert isinstance(r, Axes) |
| return |
|
|
| assert isinstance(returned, Series) |
|
|
| assert sorted(returned.keys()) == sorted(expected_keys) |
| for key, value in returned.items(): |
| assert isinstance(value, types[return_type]) |
| |
| if return_type == "axes": |
| if check_ax_title: |
| assert value.get_title() == key |
| elif return_type == "both": |
| if check_ax_title: |
| assert value.ax.get_title() == key |
| assert isinstance(value.ax, Axes) |
| assert isinstance(value.lines, dict) |
| elif return_type == "dict": |
| line = value["medians"][0] |
| axes = line.axes |
| if check_ax_title: |
| assert axes.get_title() == key |
| else: |
| raise AssertionError |
|
|
|
|
| def _check_grid_settings(obj, kinds, kws={}): |
| |
|
|
| import matplotlib as mpl |
|
|
| def is_grid_on(): |
| xticks = mpl.pyplot.gca().xaxis.get_major_ticks() |
| yticks = mpl.pyplot.gca().yaxis.get_major_ticks() |
| xoff = all(not g.gridline.get_visible() for g in xticks) |
| yoff = all(not g.gridline.get_visible() for g in yticks) |
|
|
| return not (xoff and yoff) |
|
|
| spndx = 1 |
| for kind in kinds: |
| mpl.pyplot.subplot(1, 4 * len(kinds), spndx) |
| spndx += 1 |
| mpl.rc("axes", grid=False) |
| obj.plot(kind=kind, **kws) |
| assert not is_grid_on() |
| mpl.pyplot.clf() |
|
|
| mpl.pyplot.subplot(1, 4 * len(kinds), spndx) |
| spndx += 1 |
| mpl.rc("axes", grid=True) |
| obj.plot(kind=kind, grid=False, **kws) |
| assert not is_grid_on() |
| mpl.pyplot.clf() |
|
|
| if kind not in ["pie", "hexbin", "scatter"]: |
| mpl.pyplot.subplot(1, 4 * len(kinds), spndx) |
| spndx += 1 |
| mpl.rc("axes", grid=True) |
| obj.plot(kind=kind, **kws) |
| assert is_grid_on() |
| mpl.pyplot.clf() |
|
|
| mpl.pyplot.subplot(1, 4 * len(kinds), spndx) |
| spndx += 1 |
| mpl.rc("axes", grid=False) |
| obj.plot(kind=kind, grid=True, **kws) |
| assert is_grid_on() |
| mpl.pyplot.clf() |
|
|
|
|
| def _unpack_cycler(rcParams, field="color"): |
| """ |
| Auxiliary function for correctly unpacking cycler after MPL >= 1.5 |
| """ |
| return [v[field] for v in rcParams["axes.prop_cycle"]] |
|
|
|
|
| def get_x_axis(ax): |
| return ax._shared_axes["x"] |
|
|
|
|
| def get_y_axis(ax): |
| return ax._shared_axes["y"] |
|
|
|
|
| def _check_plot_works(f, default_axes=False, **kwargs): |
| """ |
| Create plot and ensure that plot return object is valid. |
| |
| Parameters |
| ---------- |
| f : func |
| Plotting function. |
| default_axes : bool, optional |
| If False (default): |
| - If `ax` not in `kwargs`, then create subplot(211) and plot there |
| - Create new subplot(212) and plot there as well |
| - Mind special corner case for bootstrap_plot (see `_gen_two_subplots`) |
| If True: |
| - Simply run plotting function with kwargs provided |
| - All required axes instances will be created automatically |
| - It is recommended to use it when the plotting function |
| creates multiple axes itself. It helps avoid warnings like |
| 'UserWarning: To output multiple subplots, |
| the figure containing the passed axes is being cleared' |
| **kwargs |
| Keyword arguments passed to the plotting function. |
| |
| Returns |
| ------- |
| Plot object returned by the last plotting. |
| """ |
| import matplotlib.pyplot as plt |
|
|
| if default_axes: |
| gen_plots = _gen_default_plot |
| else: |
| gen_plots = _gen_two_subplots |
|
|
| ret = None |
| try: |
| fig = kwargs.get("figure", plt.gcf()) |
| plt.clf() |
|
|
| for ret in gen_plots(f, fig, **kwargs): |
| tm.assert_is_valid_plot_return_object(ret) |
|
|
| finally: |
| plt.close(fig) |
|
|
| return ret |
|
|
|
|
| def _gen_default_plot(f, fig, **kwargs): |
| """ |
| Create plot in a default way. |
| """ |
| yield f(**kwargs) |
|
|
|
|
| def _gen_two_subplots(f, fig, **kwargs): |
| """ |
| Create plot on two subplots forcefully created. |
| """ |
| if "ax" not in kwargs: |
| fig.add_subplot(211) |
| yield f(**kwargs) |
|
|
| if f is pd.plotting.bootstrap_plot: |
| assert "ax" not in kwargs |
| else: |
| kwargs["ax"] = fig.add_subplot(212) |
| yield f(**kwargs) |
|
|
