Use axes=self as an extra args in parasite_axes

weiji-li · ksunden · commit 96c61689274a · 2022-12-06T15:08:48.000-06:00
Right now parasite_axes just use self._parent_axes._get_lines as self._get_lines, but it can't update the axes unit when there are twin axes. Therefore, we need to provide axes=self as an extra args to handle this. We also need to change the callees to use axes in kwargs when provided.

Add a unit test for missing unit in twin axes

This test creates a plot with twin axes where both axes have units. It then checks whether units are appended correctly on the respective axes. The code base without the modification fails the unit test whereas the modification makes it pass the unit test.t:q`
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -239,7 +239,9 @@ def set_prop_cycle(self, cycler):
         self._prop_keys = cycler.keys  # This should make a copy
 
     def __call__(self, *args, data=None, **kwargs):
-        self.axes._process_unit_info(kwargs=kwargs)
+        # use the axes from kwargs if provided
+        axes = kwargs.get("axes", self.axes)
+        axes._process_unit_info(kwargs=kwargs)
 
         for pos_only in "xy":
             if pos_only in kwargs:
@@ -496,10 +498,12 @@ def _plot_args(self, tup, kwargs, *,
         else:
             x, y = index_of(xy[-1])
 
-        if self.axes.xaxis is not None:
-            self.axes.xaxis.update_units(x)
-        if self.axes.yaxis is not None:
-            self.axes.yaxis.update_units(y)
+        # use the axes from kwargs if provided when updating units
+        axes = kwargs.get("axes", self.axes)
+        if axes.xaxis is not None:
+            axes.xaxis.update_units(x)
+        if axes.yaxis is not None:
+            axes.yaxis.update_units(y)
 
         if x.shape[0] != y.shape[0]:
             raise ValueError(f"x and y must have same first dimension, but "
diff --git a/lib/mpl_toolkits/axes_grid1/parasite_axes.py b/lib/mpl_toolkits/axes_grid1/parasite_axes.py
@@ -3,6 +3,7 @@
 import matplotlib.transforms as mtransforms
 from matplotlib.transforms import Bbox
 from .mpl_axes import Axes
+import functools
 
 
 class ParasiteAxesBase:
@@ -18,7 +19,8 @@ def __init__(self, parent_axes, aux_transform=None,
     def clear(self):
         super().clear()
         martist.setp(self.get_children(), visible=False)
-        self._get_lines = self._parent_axes._get_lines
+        self._get_lines = functools.partial(
+            self._parent_axes._get_lines, axes=self)
         self._parent_axes.callbacks._connect_picklable(
             "xlim_changed", self._sync_lims)
         self._parent_axes.callbacks._connect_picklable(
diff --git a/lib/mpl_toolkits/axes_grid1/tests/baseline_images/test_axes_grid1/twin_axes_both_with_units.png b/lib/mpl_toolkits/axes_grid1/tests/baseline_images/test_axes_grid1/twin_axes_both_with_units.png
diff --git a/lib/mpl_toolkits/axes_grid1/tests/test_axes_grid1.py b/lib/mpl_toolkits/axes_grid1/tests/test_axes_grid1.py
@@ -25,7 +25,8 @@
     zoomed_inset_axes, mark_inset, inset_axes, BboxConnectorPatch,
     InsetPosition)
 import mpl_toolkits.axes_grid1.mpl_axes
-
+import matplotlib.units as units
+import matplotlib.ticker as ticker
 import pytest
 
 import numpy as np
@@ -87,6 +88,43 @@ def test_twin_axes_empty_and_removed():
     plt.subplots_adjust(wspace=0.5, hspace=1)
 
 
+@image_comparison(['twin_axes_both_with_units.png'])
+def test_twin_axes_both_with_units():
+    mpl.rcParams.update(
+        {"font.size": 8, "xtick.labelsize": 8, "ytick.labelsize": 8})
+    class TestUnit:
+        def __init__(self, val):
+            self._val = val
+
+    class UnitA(TestUnit):
+        fmt = "%0.1f Unit A"
+    class UnitB(TestUnit):
+        fmt = "%0.1f Unit B"
+
+    class UnitConverter(units.ConversionInterface):
+        @staticmethod
+        def convert(value, unit, axis):
+            return [x._val for x in value]
+
+        @staticmethod
+        def axisinfo(unit, axis):
+            return units.AxisInfo(majfmt=ticker.FormatStrFormatter(unit.fmt))
+
+        @staticmethod
+        def default_units(x, axis):
+            return x[0].__class__
+
+    units.registry[UnitA] = UnitConverter()
+    units.registry[UnitB] = UnitConverter()
+
+    host = host_subplot(111)
+
+    p1, = host.plot([0, 1, 2], [UnitA(x) for x in (0, 1, 2)])
+    par1 = host.twinx()
+    par1.axis["right"].major_ticklabels.set_visible(True)
+    p2, = par1.plot([0, 1, 2], [UnitB(x) for x in (0, 2, 2)])
+
+
 def test_axesgrid_colorbar_log_smoketest():
     fig = plt.figure()
     grid = AxesGrid(fig, 111,  # modified to be only subplot
