FIX: brokenbarh math before units

jklymak · jklymak · commit 78b1bfebde23 · 2019-01-04T15:27:32.000-08:00
TST: Add test of broken_barh
FIX: make bar work with timedeltas
TST: check bar and barh work with timedelta
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2011,6 +2011,19 @@ 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(x0, x, dx, convert):
+        """ Small helper to do logic of width conversion flexibly """
+        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 +2185,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(x0, x, width, self.convert_xunits)
             if xerr is not None:
-                xerr = self.convert_xunits(xerr)
+                xerr = self._convert_dx(x0, x, xerr, 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(y0, y, height, 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(y0, y, yerr, self.convert_yunits)
 
         # Now that units have been converted, set the tick locations.
         if orientation == 'vertical':
@@ -2465,10 +2480,16 @@ def broken_barh(self, xranges, yrange, **kwargs):
         self._process_unit_info(xdata=xdata,
                                 ydata=ydata,
                                 kwargs=kwargs)
-        xranges = self.convert_xunits(xranges)
+        xnew = []
+        for xr in xranges:
+            # convert the absolute values, not the x and dx...
+            x0 = self.convert_xunits(xr[0])
+            x1 = self._convert_dx(xr[0], x0, xr[1], self.convert_xunits)
+            xnew.append((x0, x1))
+
         yrange = self.convert_yunits(yrange)
 
-        col = mcoll.BrokenBarHCollection(xranges, yrange, **kwargs)
+        col = mcoll.BrokenBarHCollection(xnew, yrange, **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,29 @@ 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]])
+
+
 @image_comparison(baseline_images=['hist_log'],
                   remove_text=True)
 def test_hist_log():
@@ -5433,6 +5456,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)
