#!/usr/bin/env python # coding: utf-8 import pytest import os import warnings from pandas import DataFrame, Series from pandas.compat import zip, iteritems from pandas.util._decorators import cache_readonly from pandas.core.dtypes.api import is_list_like import pandas.util.testing as tm from pandas.util.testing import (ensure_clean, assert_is_valid_plot_return_object) import pandas.util._test_decorators as td import numpy as np from numpy import random import pandas.plotting as plotting from pandas.plotting._tools import _flatten """ This is a common base class used for various plotting tests """ def _skip_if_no_scipy_gaussian_kde(): try: from scipy.stats import gaussian_kde # noqa except ImportError: pytest.skip("scipy version doesn't support gaussian_kde") def _ok_for_gaussian_kde(kind): if kind in ['kde', 'density']: try: from scipy.stats import gaussian_kde # noqa except ImportError: return False return plotting._compat._mpl_ge_1_5_0() @td.skip_if_no_mpl class TestPlotBase(object): def setup_method(self, method): import matplotlib as mpl mpl.rcdefaults() self.mpl_le_1_2_1 = plotting._compat._mpl_le_1_2_1() self.mpl_ge_1_3_1 = plotting._compat._mpl_ge_1_3_1() self.mpl_ge_1_4_0 = plotting._compat._mpl_ge_1_4_0() self.mpl_ge_1_5_0 = plotting._compat._mpl_ge_1_5_0() self.mpl_ge_2_0_0 = plotting._compat._mpl_ge_2_0_0() self.mpl_ge_2_0_1 = plotting._compat._mpl_ge_2_0_1() self.mpl_ge_2_2_0 = plotting._compat._mpl_ge_2_2_0() if self.mpl_ge_1_4_0: self.bp_n_objects = 7 else: self.bp_n_objects = 8 if self.mpl_ge_1_5_0: # 1.5 added PolyCollections to legend handler # so we have twice as many items. self.polycollection_factor = 2 else: self.polycollection_factor = 1 if self.mpl_ge_2_0_0: self.default_figsize = (6.4, 4.8) else: self.default_figsize = (8.0, 6.0) self.default_tick_position = 'left' if self.mpl_ge_2_0_0 else 'default' n = 100 with tm.RNGContext(42): gender = np.random.choice(['Male', 'Female'], size=n) classroom = np.random.choice(['A', 'B', 'C'], size=n) self.hist_df = DataFrame({'gender': gender, 'classroom': classroom, 'height': random.normal(66, 4, size=n), 'weight': random.normal(161, 32, size=n), 'category': random.randint(4, size=n)}) self.tdf = tm.makeTimeDataFrame() self.hexbin_df = DataFrame({"A": np.random.uniform(size=20), "B": np.random.uniform(size=20), "C": np.arange(20) + np.random.uniform( size=20)}) def teardown_method(self, method): tm.close() @cache_readonly def plt(self): import matplotlib.pyplot as plt return plt @cache_readonly def colorconverter(self): import matplotlib.colors as colors return colors.colorConverter def _check_legend_labels(self, 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 = self._flatten_visible(axes) for ax in axes: if visible: assert ax.get_legend() is not None self._check_text_labels(ax.get_legend().get_texts(), labels) else: assert ax.get_legend() is None def _check_data(self, xp, rs): """ Check each axes has identical lines Parameters ---------- xp : matplotlib Axes object rs : matplotlib Axes object """ xp_lines = xp.get_lines() rs_lines = rs.get_lines() def check_line(xpl, rsl): xpdata = xpl.get_xydata() rsdata = rsl.get_xydata() tm.assert_almost_equal(xpdata, rsdata) assert len(xp_lines) == len(rs_lines) [check_line(xpl, rsl) for xpl, rsl in zip(xp_lines, rs_lines)] tm.close() def _check_visible(self, 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 _get_colors_mapped(self, series, colors): unique = series.unique() # unique and colors length can be differed # depending on slice value mapped = dict(zip(unique, colors)) return [mapped[v] for v in series.values] def _check_colors(self, 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.lines import Line2D from matplotlib.collections import ( Collection, PolyCollection, LineCollection ) conv = self.colorconverter if linecolors is not None: if mapping is not None: linecolors = self._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() # Line2D may contains string color expression 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 = self._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): # returned as list of np.array 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(self, 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 l, e in zip(labels, expected): assert l == e def _check_ticks_props(self, 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 = self._flatten_visible(axes) for ax in axes: if xlabelsize or xrot: if isinstance(ax.xaxis.get_minor_formatter(), NullFormatter): # If minor ticks has NullFormatter, rot / fontsize are not # retained 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 or yrot: 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(self, 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 = self._flatten_visible(axes) for ax in axes: assert ax.xaxis.get_scale() == xaxis assert ax.yaxis.get_scale() == yaxis def _check_axes_shape(self, 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 """ if figsize is None: figsize = self.default_figsize visible_axes = self._flatten_visible(axes) if axes_num is not None: assert len(visible_axes) == axes_num for ax in visible_axes: # check something drawn on visible axes assert len(ax.get_children()) > 0 if layout is not None: result = self._get_axes_layout(_flatten(axes)) assert result == layout tm.assert_numpy_array_equal( visible_axes[0].figure.get_size_inches(), np.array(figsize, dtype=np.float64)) def _get_axes_layout(self, axes): x_set = set() y_set = set() for ax in axes: # check axes coordinates to estimate layout points = ax.get_position().get_points() x_set.add(points[0][0]) y_set.add(points[0][1]) return (len(y_set), len(x_set)) def _flatten_visible(self, axes): """ Flatten axes, and filter only visible Parameters ---------- axes : matplotlib Axes object, or its list-like """ axes = _flatten(axes) axes = [ax for ax in axes if ax.get_visible()] return axes def _check_has_errorbars(self, 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 = self._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(self, 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: # should be fixed when the returning default is changed 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: # should be fixed when the returning default is changed if return_type is None: for r in self._flatten_visible(returned): assert isinstance(r, Axes) return assert isinstance(returned, Series) assert sorted(returned.keys()) == sorted(expected_keys) for key, value in iteritems(returned): assert isinstance(value, types[return_type]) # check returned dict has correct mapping 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 self.mpl_ge_1_5_0 else line.get_axes() if check_ax_title: assert axes.get_title() == key else: raise AssertionError def _check_grid_settings(self, obj, kinds, kws={}): # Make sure plot defaults to rcParams['axes.grid'] setting, GH 9792 import matplotlib as mpl def is_grid_on(): xoff = all(not g.gridOn for g in self.plt.gca().xaxis.get_major_ticks()) yoff = all(not g.gridOn for g in self.plt.gca().yaxis.get_major_ticks()) return not (xoff and yoff) spndx = 1 for kind in kinds: if not _ok_for_gaussian_kde(kind): continue self.plt.subplot(1, 4 * len(kinds), spndx) spndx += 1 mpl.rc('axes', grid=False) obj.plot(kind=kind, **kws) assert not is_grid_on() self.plt.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() if kind != 'pie': self.plt.subplot(1, 4 * len(kinds), spndx) spndx += 1 mpl.rc('axes', grid=True) obj.plot(kind=kind, **kws) assert is_grid_on() self.plt.subplot(1, 4 * len(kinds), spndx) spndx += 1 mpl.rc('axes', grid=False) obj.plot(kind=kind, grid=True, **kws) assert is_grid_on() def _maybe_unpack_cycler(self, rcParams, field='color'): """ Compat layer for MPL 1.5 change to color cycle Before: plt.rcParams['axes.color_cycle'] -> ['b', 'g', 'r'...] After : plt.rcParams['axes.prop_cycle'] -> cycler(...) """ if self.mpl_ge_1_5_0: cyl = rcParams['axes.prop_cycle'] colors = [v[field] for v in cyl] else: colors = rcParams['axes.color_cycle'] return colors def _check_plot_works(f, filterwarnings='always', **kwargs): import matplotlib.pyplot as plt ret = None with warnings.catch_warnings(): warnings.simplefilter(filterwarnings) try: try: fig = kwargs['figure'] except KeyError: fig = plt.gcf() plt.clf() ax = kwargs.get('ax', fig.add_subplot(211)) # noqa ret = f(**kwargs) assert_is_valid_plot_return_object(ret) try: kwargs['ax'] = fig.add_subplot(212) ret = f(**kwargs) except Exception: pass else: assert_is_valid_plot_return_object(ret) with ensure_clean(return_filelike=True) as path: plt.savefig(path) finally: tm.close(fig) return ret def curpath(): pth, _ = os.path.split(os.path.abspath(__file__)) return pth