Merge pull request #19089 from jklymak/enh-pass-scale-to-scale

timhoffm · web-flow · commit a4cc0fc60493 · 2020-12-11T01:41:26.000+01:00
ENH: allow passing a scale instance to set_scale
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -3510,7 +3510,7 @@ def set_xscale(self, value, **kwargs):
 
         Parameters
         ----------
-        value : {"linear", "log", "symlog", "logit", ...}
+        value : {"linear", "log", "symlog", "logit", ...} or `.ScaleBase`
             The axis scale type to apply.
 
         **kwargs
@@ -3836,7 +3836,7 @@ def set_yscale(self, value, **kwargs):
 
         Parameters
         ----------
-        value : {"linear", "log", "symlog", "logit", ...}
+        value : {"linear", "log", "symlog", "logit", ...} or `.ScaleBase`
             The axis scale type to apply.
 
         **kwargs
diff --git a/lib/matplotlib/axis.py b/lib/matplotlib/axis.py
@@ -741,7 +741,10 @@ def get_scale(self):
         return self._scale.name
 
     def _set_scale(self, value, **kwargs):
-        self._scale = mscale.scale_factory(value, self, **kwargs)
+        if not isinstance(value, mscale.ScaleBase):
+            self._scale = mscale.scale_factory(value, self, **kwargs)
+        else:
+            self._scale = value
         self._scale.set_default_locators_and_formatters(self)
 
         self.isDefault_majloc = True
diff --git a/lib/matplotlib/tests/test_scale.py b/lib/matplotlib/tests/test_scale.py
@@ -2,6 +2,7 @@
 from matplotlib.scale import (
     LogTransform, InvertedLogTransform,
     SymmetricalLogTransform)
+import matplotlib.scale as mscale
 from matplotlib.testing.decorators import check_figures_equal, image_comparison
 
 import numpy as np
@@ -198,3 +199,14 @@ def forward(x):
     ax.plot(x, x)
     ax.set_xscale('function', functions=(forward, inverse))
     ax.set_xlim(1, 1000)
+
+
+def test_pass_scale():
+    # test passing a scale object works...
+    fig, ax = plt.subplots()
+    scale = mscale.LogScale(axis=None)
+    ax.set_xscale(scale)
+    scale = mscale.LogScale(axis=None)
+    ax.set_yscale(scale)
+    assert ax.xaxis.get_scale() == 'log'
+    assert ax.yaxis.get_scale() == 'log'
