BUG/CLN: Refactoring tsplot

Author: sinhrks

pandas/tools/plotting.py

@@ -1552,18 +1552,6 @@ def __init__(self, data, **kwargs):
         if 'x_compat' in self.kwds:
             self.x_compat = bool(self.kwds.pop('x_compat'))
 
-    def _index_freq(self):
-        from pandas.core.frame import DataFrame
-        if isinstance(self.data, (Series, DataFrame)):
-            freq = getattr(self.data.index, 'freq', None)
-            if freq is None:
-                freq = getattr(self.data.index, 'inferred_freq', None)
-                if freq == 'B':
-                    weekdays = np.unique(self.data.index.dayofweek)
-                    if (5 in weekdays) or (6 in weekdays):
-                        freq = None
-            return freq
-
     def _is_dynamic_freq(self, freq):
         if isinstance(freq, DateOffset):
             freq = freq.rule_code
@@ -1574,9 +1562,7 @@ def _is_dynamic_freq(self, freq):
 
     def _no_base(self, freq):
         # hack this for 0.10.1, creating more technical debt...sigh
-        from pandas.core.frame import DataFrame
-        if (isinstance(self.data, (Series, DataFrame))
-            and isinstance(self.data.index, DatetimeIndex)):
+        if isinstance(self.data.index, DatetimeIndex):
             base = frequencies.get_freq(freq)
             x = self.data.index
             if (base <= frequencies.FreqGroup.FR_DAY):
@@ -1586,7 +1572,8 @@ def _no_base(self, freq):
         return True
 
     def _use_dynamic_x(self):
-        freq = self._index_freq()
+        from pandas.tseries.plotting import _get_index_freq
+        freq = _get_index_freq(self.data)
 
         ax = self._get_ax(0)
         ax_freq = getattr(ax, 'freq', None)
@@ -1607,10 +1594,22 @@ def _make_plot(self):
 
         if self._is_ts_plot():
             data = self._maybe_convert_index(self.data)
+            freq = data.index.freq
+            kwds = self.kwds.copy()
+            from pandas.tseries.plotting import _maybe_resample
+            for ax in self.axes:
+                # resample data and replot if required
+                data = _maybe_resample(data, ax, freq, kwds)
+
             x = data.index      # dummy, not used
             plotf = self._get_ts_plot_function()
             it = self._iter_data(data=data, keep_index=True)
         else:
+            from pandas.tseries.plotting import _replot_x_compat
+            for ax in self.axes:
+                # if ax holds _plot_data, replot them on the x_compat scale
+                _replot_x_compat(ax)
+
             x = self._get_xticks(convert_period=True)
             plotf = self._get_plot_function()
             it = self._iter_data()
@@ -1684,30 +1683,13 @@ def _update_prior(self, y):
                 self._neg_prior += y
 
     def _maybe_convert_index(self, data):
-        # tsplot converts automatically, but don't want to convert index
-        # over and over for DataFrames
-        from pandas.core.frame import DataFrame
-        if (isinstance(data.index, DatetimeIndex) and
-                isinstance(data, DataFrame)):
-            freq = getattr(data.index, 'freq', None)
-
-            if freq is None:
-                freq = getattr(data.index, 'inferred_freq', None)
-            if isinstance(freq, DateOffset):
-                freq = freq.rule_code
-            freq = frequencies.get_base_alias(freq)
-            freq = frequencies.get_period_alias(freq)
-
-            if freq is None:
-                ax = self._get_ax(0)
-                freq = getattr(ax, 'freq', None)
-
+        # tsplot only accepts PeriodIndex
+        if isinstance(data.index, DatetimeIndex):
+            from pandas.tseries.plotting import _get_freq
+            freq = _get_freq(self._get_ax(0), self.data)
             if freq is None:
                 raise ValueError('Could not get frequency alias for plotting')
-
-            data = DataFrame(data.values,
-                             index=data.index.to_period(freq=freq),
-                             columns=data.columns)
+            data = data.to_period(freq=freq)
         return data
 
     def _post_plot_logic(self):

pandas/tseries/plotting.py

@@ -8,21 +8,18 @@
 
 import numpy as np
 
-from pandas import isnull
-from pandas.tseries.period import Period
+from pandas.tseries.period import Period, PeriodIndex
 from pandas.tseries.offsets import DateOffset
 import pandas.tseries.frequencies as frequencies
-from pandas.tseries.index import DatetimeIndex
-import pandas.core.common as com
 
-from pandas.tseries.converter import (PeriodConverter, TimeSeries_DateLocator,
+from pandas.tseries.converter import (TimeSeries_DateLocator,
                                       TimeSeries_DateFormatter)
 
 #----------------------------------------------------------------------
 # Plotting functions and monkey patches
 
 
-def tsplot(series, plotf, **kwargs):
+def tsplot(series, plotf, ax=None, **kwargs):
     """
     Plots a Series on the given Matplotlib axes or the current axes
 
@@ -36,36 +33,21 @@ def tsplot(series, plotf, **kwargs):
     Supports same kwargs as Axes.plot
 
     """
-    # Used inferred freq is possible, need a test case for inferred
-    if 'ax' in kwargs:
-        ax = kwargs.pop('ax')
-    else:
-        import matplotlib.pyplot as plt
-        ax = plt.gca()
-
-    freq = _get_freq(ax, series)
-    # resample against axes freq if necessary
-    if freq is None:  # pragma: no cover
-        raise ValueError('Cannot use dynamic axis without frequency info')
-    else:
-        # Convert DatetimeIndex to PeriodIndex
-        if isinstance(series.index, DatetimeIndex):
-            series = series.to_period(freq=freq)
-        freq, ax_freq, series = _maybe_resample(series, ax, freq, plotf,
-                                                kwargs)
-
+    if ax is None:
+        raise ValueError('ax keyword must be passed')
+    if not isinstance(series.index, PeriodIndex):
+        raise ValueError('Passed data must have PeriodIndex')
     # Set ax with freq info
-    _decorate_axes(ax, freq, kwargs)
+    _decorate_axes(ax, series.index.freq)
 
     # how to make sure ax.clear() flows through?
-    if not hasattr(ax, '_plot_data'):
+    if getattr(ax, '_plot_data', None) is None:
         ax._plot_data = []
     ax._plot_data.append((series, plotf, kwargs))
     lines = plotf(ax, series.index._mpl_repr(), series.values, **kwargs)
 
     # set date formatter, locators and rescale limits
     format_dateaxis(ax, ax.freq)
-
     # x and y coord info
     ax.format_coord = lambda t, y: ("t = {0}  "
                                     "y = {1:8f}".format(Period(ordinal=int(t),
@@ -75,24 +57,24 @@ def tsplot(series, plotf, **kwargs):
     return lines
 
 
-def _maybe_resample(series, ax, freq, plotf, kwargs):
+def _maybe_resample(data, ax, freq, kwargs):
     ax_freq = _get_ax_freq(ax)
     if ax_freq is not None and freq != ax_freq:
         if frequencies.is_superperiod(freq, ax_freq):  # upsample input
-            series = series.copy()
-            series.index = series.index.asfreq(ax_freq, how='s')
+            data = data.copy()
+            data.index = data.index.asfreq(ax_freq, how='s')
             freq = ax_freq
         elif _is_sup(freq, ax_freq):  # one is weekly
             how = kwargs.pop('how', 'last')
-            series = series.resample('D', how=how).dropna()
-            series = series.resample(ax_freq, how=how).dropna()
+            data = data.resample('D', how=how).dropna()
+            data = data.resample(ax_freq, how=how).dropna()
             freq = ax_freq
         elif frequencies.is_subperiod(freq, ax_freq) or _is_sub(freq, ax_freq):
-            _upsample_others(ax, freq, plotf, kwargs)
+            _upsample_others(ax, freq)
             ax_freq = freq
         else:  # pragma: no cover
             raise ValueError('Incompatible frequency conversion')
-    return freq, ax_freq, series
+    return data
 
 
 def _get_ax_freq(ax):
@@ -115,67 +97,91 @@ def _is_sup(f1, f2):
             (f2.startswith('W') and frequencies.is_superperiod(f1, 'D')))
 
 
-def _upsample_others(ax, freq, plotf, kwargs):
-    legend = ax.get_legend()
-    lines, labels = _replot_ax(ax, freq, kwargs)
+def _upsample_others(ax, freq):
+
+    def _replot(ax):
+        data = getattr(ax, '_plot_data', None)
+        if data is None:
+            return
+
+        # preserve legend
+        leg = ax.get_legend()
+        handles, labels = ax.get_legend_handles_labels()
+
+        ax._plot_data = []
+        ax.clear()
+        _decorate_axes(ax, freq)
 
-    other_ax = None
-    if hasattr(ax, 'left_ax'):
-        other_ax = ax.left_ax
-    if hasattr(ax, 'right_ax'):
-        other_ax = ax.right_ax
-
-    if other_ax is not None:
-        rlines, rlabels = _replot_ax(other_ax, freq, kwargs)
-        lines.extend(rlines)
-        labels.extend(rlabels)
-
-    if (legend is not None and kwargs.get('legend', True) and
-            len(lines) > 0):
-        title = legend.get_title().get_text()
-        if title == 'None':
-            title = None
-        ax.legend(lines, labels, loc='best', title=title)
-
-
-def _replot_ax(ax, freq, kwargs):
-    data = getattr(ax, '_plot_data', None)
-    ax._plot_data = []
-    ax.clear()
-    _decorate_axes(ax, freq, kwargs)
-
-    lines = []
-    labels = []
-    if data is not None:
         for series, plotf, kwds in data:
             series = series.copy()
             idx = series.index.asfreq(freq, how='S')
             series.index = idx
-            ax._plot_data.append(series)
-            lines.append(plotf(ax, series.index._mpl_repr(), series.values, **kwds)[0])
-            labels.append(com.pprint_thing(series.name))
+            ax._plot_data.append((series, plotf, kwds))
+            plotf(ax, series.index._mpl_repr(), series, **kwds)
+
+        if leg is not None:
+            ax.legend(handles, labels, title=leg.get_title().get_text())
+
+    _replot(ax)
+    if hasattr(ax, 'left_ax'):
+        _replot(ax.left_ax)
+    elif hasattr(ax, 'right_ax'):
+        _replot(ax.right_ax)
+
+
+def _replot_x_compat(ax):
+
+    def _replot(ax):
+        data = getattr(ax, '_plot_data', None)
+        if data is None:
+            return
+
+        # preserve legend
+        leg = ax.get_legend()
+        handles, labels = ax.get_legend_handles_labels()
+
+        ax._plot_data = None
+        ax.clear()
+
+        _decorate_axes(ax, None)
+
+        for series, plotf, kwds in data:
+            idx = series.index.to_timestamp(how='e')
+            series.index = idx
+            plotf(ax, series.index_mpl_repr(), series, **kwds)
 
-    return lines, labels
+        if leg is not None:
+            ax.legend(handles, labels, title=leg.get_title().get_text())
+
+    _replot(ax)
+    if hasattr(ax, 'left_ax'):
+        _replot(ax.left_ax)
+    elif hasattr(ax, 'right_ax'):
+        _replot(ax.right_ax)
 
 
-def _decorate_axes(ax, freq, kwargs):
+def _decorate_axes(ax, freq):
     ax.freq = freq
     xaxis = ax.get_xaxis()
     xaxis.freq = freq
-    if not hasattr(ax, 'legendlabels'):
-        ax.legendlabels = [kwargs.get('label', None)]
-    else:
-        ax.legendlabels.append(kwargs.get('label', None))
     ax.view_interval = None
     ax.date_axis_info = None
 
 
-def _get_freq(ax, series):
-    # get frequency from data
-    freq = getattr(series.index, 'freq', None)
+def _get_index_freq(data):
+    freq = getattr(data.index, 'freq', None)
     if freq is None:
-        freq = getattr(series.index, 'inferred_freq', None)
+        freq = getattr(data.index, 'inferred_freq', None)
+        if freq == 'B':
+            weekdays = np.unique(data.index.dayofweek)
+            if (5 in weekdays) or (6 in weekdays):
+                freq = None
+    return freq
+
 
+def _get_freq(ax, data):
+    # get frequency from data
+    freq = _get_index_freq(data)
     ax_freq = getattr(ax, 'freq', None)
 
     # use axes freq if no data freq

pandas/tseries/tests/test_plotting.py

@@ -91,19 +91,16 @@ def test_nonnumeric_exclude(self):
 
     @slow
     def test_tsplot(self):
-        from pandas.tseries.plotting import tsplot
         import matplotlib.pyplot as plt
 
         ax = plt.gca()
         ts = tm.makeTimeSeries()
 
-        f = lambda *args, **kwds: tsplot(s, plt.Axes.plot, *args, **kwds)
-
         for s in self.period_ser:
-            _check_plot_works(f, s.index.freq, ax=ax, series=s)
+            _check_plot_works(s.plot, ax=ax)
 
         for s in self.datetime_ser:
-            _check_plot_works(f, s.index.freq.rule_code, ax=ax, series=s)
+            _check_plot_works(s.plot, ax=ax)
 
         ax = ts.plot(style='k')
         self.assertEqual((0., 0., 0.), ax.get_lines()[0].get_color())