Add test for axis label when units set

dstansby · dstansby · commit f7c4bb02cdfa · 2022-07-12T16:01:58.000+01:00
Fix indentation
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/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -283,3 +283,23 @@ def test_plot_kernel():
     # just a smoketest that fail
     kernel = Kernel([1, 2, 3, 4, 5])
     plt.plot(kernel)
+
+
+@check_figures_equal(extensions=['png'])
+def test_unit_axis_label(fig_test, fig_ref):
+    import matplotlib.testing.jpl_units as units
+    units.register()
+
+    data = [0 * units.km, 1 * units.km, 2 * units.km]
+
+    ax_test = fig_test.subplots()
+    ax_ref = fig_ref.subplots()
+    axs = [ax_test, ax_ref]
+
+    for ax in axs:
+        ax.yaxis.set_units('km')
+        ax.set_xlim(10, 20)
+        ax.set_ylim(10, 20)
+
+    ax_test.scatter([1, 2, 3], data, edgecolors='none')
+    ax_ref.plot([1, 2, 3], data, marker='o', linewidth=0)
