Merge pull request #12903 from jklymak/fix-broken-barh-units

anntzer · web-flow · commit b7b8f5b2e9a5 · 2019-01-13T12:08:49.000+01:00
FIX: (broken)bar(h) math before units
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2011,6 +2011,58 @@ 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, xconv, convert):
+        """
+        Small helper to do logic of width conversion flexibly.
+
+        *dx* and *x0* have units, but *xconv* has already been converted
+        to unitless (and is an ndarray).  This allows the *dx* to have units
+        that are different from *x0*, but are still accepted by the
+        ``__add__`` operator of *x0*.
+        """
+
+        # x should be an array...
+        assert type(xconv) is np.ndarray
+
+        if xconv.size == 0:
+            # xconv has already been converted, but maybe empty...
+            return convert(dx)
+
+        try:
+            # attempt to add the width to x0; this works for
+            # datetime+timedelta, for instance
+
+            # only use the first element of x and x0.  This saves
+            # having to be sure addition works across the whole
+            # vector.  This is particularly an issue if
+            # x0 and dx are lists so x0 + dx just concatenates the lists.
+            # We can't just cast x0 and dx to numpy arrays because that
+            # removes the units from unit packages like `pint` that
+            # wrap numpy arrays.
+            try:
+                x0 = x0[0]
+            except (TypeError, IndexError, KeyError):
+                x0 = x0
+
+            try:
+                x = xconv[0]
+            except (TypeError, IndexError, KeyError):
+                x = xconv
+
+            delist = False
+            if not np.iterable(dx):
+                dx = [dx]
+                delist = True
+            dx = [convert(x0 + ddx) - x for ddx in dx]
+            if delist:
+                dx = dx[0]
+        except (TypeError, AttributeError) as e:
+            # 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",
@@ -2175,16 +2227,17 @@ def bar(self, x, height, width=0.8, bottom=None, *, align="center",
         # lets do some conversions now since some types cannot be
         # subtracted uniformly
         if self.xaxis is not None:
-            x = self.convert_xunits(x)
-            width = self.convert_xunits(width)
+            x0 = x
+            x = np.asarray(self.convert_xunits(x))
+            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:
-            y = self.convert_yunits(y)
-            height = self.convert_yunits(height)
+            y0 = y
+            y = np.asarray(self.convert_yunits(y))
+            height = self._convert_dx(height, y0, y, self.convert_yunits)
             if yerr is not None:
-                yerr = self.convert_yunits(yerr)
+                yerr = self._convert_dx(yerr, y0, y, self.convert_yunits)
 
         x, height, width, y, linewidth = np.broadcast_arrays(
             # Make args iterable too.
@@ -2465,10 +2518,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 = np.asarray(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():
@@ -5432,6 +5458,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)
