Merge pull request #23417 from dstansby/unit-axis-label

QuLogic · web-flow · commit 4e059bca2a77 · 2022-11-26T04:35:43.000-05:00
Consistently set label on axis with units
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2629,7 +2629,9 @@ def _process_unit_info(self, datasets=None, kwargs=None, *, convert=True):
             except KeyError:
                 raise ValueError(f"Invalid axis name: {axis_name!r}") from None
             # Update from data if axis is already set but no unit is set yet.
-            if axis is not None and data is not None and not axis.have_units():
+            if (axis is not None and
+                    data is not None and
+                    not axis._have_units_and_converter()):
                 axis.update_units(data)
         for axis_name, axis in axis_map.items():
             # Return if no axis is set.
diff --git a/lib/matplotlib/axis.py b/lib/matplotlib/axis.py
@@ -1715,8 +1715,17 @@ def _update_axisinfo(self):
         self.set_default_intervals()
 
     def have_units(self):
+        """
+        Return `True` if units or a converter have been set.
+        """
         return self.converter is not None or self.units is not None
 
+    def _have_units_and_converter(self):
+        """
+        Return `True` if units and a converter have been set.
+        """
+        return self.converter is not None and self.units is not None
+
     def convert_units(self, x):
         # If x is natively supported by Matplotlib, doesn't need converting
         if munits._is_natively_supported(x):
diff --git a/lib/matplotlib/testing/jpl_units/UnitDbl.py b/lib/matplotlib/testing/jpl_units/UnitDbl.py
@@ -91,6 +91,8 @@ def __bool__(self):
 
     def _cmp(self, op, rhs):
         """Check that *self* and *rhs* share units; compare them using *op*."""
+        if rhs is None:
+            return NotImplemented
         self.checkSameUnits(rhs, "compare")
         return op(self._value, rhs._value)
 
diff --git a/lib/matplotlib/tests/baseline_images/test_patches/units_rectangle.png b/lib/matplotlib/tests/baseline_images/test_patches/units_rectangle.png
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -283,3 +283,19 @@ def test_plot_kernel():
     # just a smoketest that fail
     kernel = Kernel([1, 2, 3, 4, 5])
     plt.plot(kernel)
+
+
+@pytest.mark.parametrize('plot_meth_name', ['scatter', 'plot'])
+def test_unit_axis_label(plot_meth_name):
+    # Check that the correct Axis labels are set on plots with units
+    import matplotlib.testing.jpl_units as units
+    units.register()
+
+    fig, ax = plt.subplots()
+    ax.xaxis.set_units('m')
+    ax.yaxis.set_units('sec')
+    plot_method = getattr(ax, plot_meth_name)
+    plot_method(np.arange(3) * units.m, np.arange(3) * units.sec)
+    assert ax.get_xlabel() == 'm'
+    assert ax.get_ylabel() == 'sec'
+    plt.close('all')
