FIX: brokenbarh math before units

jklymak · jklymak · commit 1dfdc298f244 · 2019-01-10T09:21:00.000-08:00
TST: Add test of broken_barh
FIX: make bar work with timedeltas
TST: check bar and barh work with timedelta
DOC: fix private method docstring
FIX: revert list comprehension

TST: fix test
FIX: throw error xr not tw-tuple
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2011,6 +2011,26 @@ def step(self, x, y, *args, where='pre', data=None, **kwargs):
         kwargs['drawstyle'] = 'steps-' + where
         return self.plot(x, y, *args, data=data, **kwargs)
 
+    @staticmethod
+    def _convert_dx(dx, x0, x, convert):
+        """
+        Small helper to do logic of width conversion flexibly.
+
+        *dx* and *x0* have units, but *x* has already been converted
+        to unitless.  This allows the *dx* to have units that are
+        different from *x0*, but are still accepted by the  ``__add__``
+        operator of *x0*.
+        """
+        try:
+            # attempt to add the width to x0; this works for
+            # datetime+timedelta, for instance
+            dx = convert(x0 + dx) - x
+        except (TypeError, AttributeError):
+            # but doesn't work for 'string' + float, so just
+            # see if the converter works on the float.
+            dx = convert(dx)
+        return dx
+
     @_preprocess_data()
     @docstring.dedent_interpd
     def bar(self, x, height, width=0.8, bottom=None, *, align="center",
@@ -2172,23 +2192,25 @@ def bar(self, x, height, width=0.8, bottom=None, *, align="center",
         else:
             raise ValueError('invalid orientation: %s' % orientation)
 
+        x, height, width, y, linewidth = np.broadcast_arrays(
+            # Make args iterable too.
+            np.atleast_1d(x), height, width, y, linewidth)
+
         # lets do some conversions now since some types cannot be
         # subtracted uniformly
         if self.xaxis is not None:
+            x0 = x
             x = self.convert_xunits(x)
-            width = self.convert_xunits(width)
+            width = self._convert_dx(width, x0, x, self.convert_xunits)
             if xerr is not None:
-                xerr = self.convert_xunits(xerr)
+                xerr = self._convert_dx(xerr, x0, x, self.convert_xunits)
 
         if self.yaxis is not None:
+            y0 = y
             y = self.convert_yunits(y)
-            height = self.convert_yunits(height)
+            height = self._convert_dx(height, y0, y, self.convert_yunits)
             if yerr is not None:
-                yerr = self.convert_yunits(yerr)
-
-        x, height, width, y, linewidth = np.broadcast_arrays(
-            # Make args iterable too.
-            np.atleast_1d(x), height, width, y, linewidth)
+                yerr = self._convert_dx(yerr, y0, y, self.convert_yunits)
 
         # Now that units have been converted, set the tick locations.
         if orientation == 'vertical':
@@ -2465,10 +2487,19 @@ def broken_barh(self, xranges, yrange, **kwargs):
         self._process_unit_info(xdata=xdata,
                                 ydata=ydata,
                                 kwargs=kwargs)
-        xranges = self.convert_xunits(xranges)
-        yrange = self.convert_yunits(yrange)
-
-        col = mcoll.BrokenBarHCollection(xranges, yrange, **kwargs)
+        xranges_conv = []
+        for xr in xranges:
+            if len(xr) != 2:
+                raise ValueError('each range in xrange must be a sequence '
+                                 'with two elements (i.e. an Nx2 array)')
+            # convert the absolute values, not the x and dx...
+            x_conv = self.convert_xunits(xr[0])
+            x1 = self._convert_dx(xr[1], xr[0], x_conv, self.convert_xunits)
+            xranges_conv.append((x_conv, x1))
+
+        yrange_conv = self.convert_yunits(yrange)
+
+        col = mcoll.BrokenBarHCollection(xranges_conv, yrange_conv, **kwargs)
         self.add_collection(col, autolim=True)
         self.autoscale_view()
 
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1498,6 +1498,32 @@ def test_barh_tick_label():
             align='center')
 
 
+def test_bar_timedelta():
+    """smoketest that bar can handle width and height in delta units"""
+    fig, ax = plt.subplots()
+    ax.bar(datetime.datetime(2018, 1, 1), 1.,
+           width=datetime.timedelta(hours=3))
+    ax.bar(datetime.datetime(2018, 1, 1), 1.,
+           xerr=datetime.timedelta(hours=2),
+           width=datetime.timedelta(hours=3))
+    fig, ax = plt.subplots()
+    ax.barh(datetime.datetime(2018, 1, 1), 1,
+            height=datetime.timedelta(hours=3))
+    ax.barh(datetime.datetime(2018, 1, 1), 1,
+            height=datetime.timedelta(hours=3),
+            yerr=datetime.timedelta(hours=2))
+    fig, ax = plt.subplots()
+    ax.barh([datetime.datetime(2018, 1, 1), datetime.datetime(2018, 1, 1)],
+            np.array([1, 1.5]),
+            height=datetime.timedelta(hours=3))
+    ax.barh([datetime.datetime(2018, 1, 1), datetime.datetime(2018, 1, 1)],
+            np.array([1, 1.5]),
+            height=[datetime.timedelta(hours=t) for t in [1, 2]])
+    ax.broken_barh([(datetime.datetime(2018, 1, 1),
+                     datetime.timedelta(hours=1))],
+                   (10, 20))
+
+
 @image_comparison(baseline_images=['hist_log'],
                   remove_text=True)
 def test_hist_log():
@@ -5433,6 +5459,15 @@ def test_broken_barh_empty():
     ax.broken_barh([], (.1, .5))
 
 
+def test_broken_barh_timedelta():
+    """Check that timedelta works as x, dx pair for this method """
+    fig, ax = plt.subplots()
+    pp = ax.broken_barh([(datetime.datetime(2018, 11, 9, 0, 0, 0),
+                          datetime.timedelta(hours=1))], [1, 2])
+    assert pp.get_paths()[0].vertices[0, 0] == 737007.0
+    assert pp.get_paths()[0].vertices[2, 0] == 737007.0 + 1 / 24
+
+
 def test_pandas_pcolormesh(pd):
     time = pd.date_range('2000-01-01', periods=10)
     depth = np.arange(20)
