Synchronize units change in Axis.set_units for shared axis

l-johnston · l-johnston · commit 6d54d3a2d97f · 2020-08-21T07:32:15.000-05:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -4305,37 +4305,3 @@ 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/axis.py b/lib/matplotlib/axis.py
@@ -1530,14 +1530,31 @@ def set_units(self, u):
         Parameters
         ----------
         u : units tag
+
+        Notes
+        -----
+        The units of any shared axis will also be updated.
         """
         if u == self.units:
             return
-        self.units = u
-        self._update_axisinfo()
-        self.callbacks.process('units')
-        self.callbacks.process('units finalize')
-        self.stale = True
+        if self is self.axes.xaxis:
+            shared = [
+                ax.xaxis
+                for ax in self.axes.get_shared_x_axes().get_siblings(self.axes)
+            ]
+        elif self is self.axes.yaxis:
+            shared = [
+                ax.yaxis
+                for ax in self.axes.get_shared_y_axes().get_siblings(self.axes)
+            ]
+        else:
+            shared = [self]
+        for axis in shared:
+            axis.units = u
+            axis._update_axisinfo()
+            axis.callbacks.process('units')
+            axis.callbacks.process('units finalize')
+            axis.stale = True
 
     def get_units(self):
         """Return the units for axis."""
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -184,16 +184,7 @@ def test_set_xyunits(quantity_converter):
     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")
+    ax1.xaxis.set_units("seconds")
+    ax2.yaxis.set_units("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"
