Add set_xunits and set_yunits

l-johnston · l-johnston · commit f77c0a57fbeb · 2020-08-20T10:47:09.000-05:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -4305,3 +4305,37 @@ def get_shared_x_axes(self):
     def get_shared_y_axes(self):
         """Return a reference to the shared axes Grouper object for y axes."""
         return self._shared_y_axes
+
+    def set_xunits(self, units, emit=True):
+        """
+        Set the x-axis units.
+
+        Parameters
+        ----------
+        units : units tag
+
+        emit : bool, default: True
+            Whether to notify observers of units change.
+        """
+        if emit:
+            for ax in self._shared_x_axes.get_siblings(self):
+                ax.xaxis.set_units(units)
+        else:
+            self.xaxis.set_units(units)
+
+    def set_yunits(self, units, emit=True):
+        """
+        Set the y-axis units.
+
+        Parameters
+        ----------
+        units : units tag
+
+        emit : bool, default: True
+            Whether to notify observers of units change.
+        """
+        if emit:
+            for ax in self._shared_y_axes.get_siblings(self):
+                ax.yaxis.set_units(units)
+        else:
+            self.yaxis.set_units(units)
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -172,3 +172,28 @@ class subdate(datetime):
 
     fig_test.subplots().plot(subdate(2000, 1, 1), 0, "o")
     fig_ref.subplots().plot(datetime(2000, 1, 1), 0, "o")
+
+
+def test_set_xyunits(quantity_converter):
+    munits.registry[Quantity] = quantity_converter
+    x = Quantity(np.linspace(0, 1, 10), "hours")
+    y1 = Quantity(np.linspace(1, 2, 10), "feet")
+    y2 = Quantity(np.linspace(3, 4, 10), "feet")
+    fig, (ax1, ax2) = plt.subplots(2, 1, sharex='all', sharey='all')
+    ax1.plot(x, y1)
+    ax2.plot(x, y2)
+    assert ax1.xaxis.get_units() == ax2.xaxis.get_units() == "hours"
+    assert ax2.yaxis.get_units() == ax2.yaxis.get_units() == "feet"
+    ax1.set_xunits("seconds")
+    ax2.set_yunits("inches")
+    assert ax1.xaxis.get_units() == ax2.xaxis.get_units() == "seconds"
+    assert ax1.yaxis.get_units() == ax2.yaxis.get_units() == "inches"
+    fig, (ax1, ax2) = plt.subplots(2, 1, sharex='all', sharey='all')
+    ax1.plot(x, y1)
+    ax2.plot(x, y2)
+    ax1.set_xunits("seconds", emit=False)
+    ax2.set_yunits("inches", emit=False)
+    assert ax1.xaxis.get_units() == "seconds"
+    assert ax2.xaxis.get_units() == "hours"
+    assert ax1.yaxis.get_units() == "feet"
+    assert ax2.yaxis.get_units() == "inches"
