Support full-sharex/y in subplot_mosaic.

anntzer · anntzer · commit 1fd0914aefb7 · 2021-05-01T12:27:02.000+02:00
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -807,8 +807,7 @@ def subplots(self, nrows=1, ncols=1, sharex=False, sharey=False,
             Number of rows/columns of the subplot grid.
 
         sharex, sharey : bool or {'none', 'all', 'row', 'col'}, default: False
-            Controls sharing of properties among x (*sharex*) or y (*sharey*)
-            axes:
+            Controls sharing of x-axis (*sharex*) or y-axis (*sharey*):
 
             - True or 'all': x- or y-axis will be shared among all subplots.
             - False or 'none': each subplot x- or y-axis will be independent.
@@ -1672,8 +1671,8 @@ def _normalize_grid_string(layout):
             layout = inspect.cleandoc(layout)
             return [list(ln) for ln in layout.strip('\n').split('\n')]
 
-    def subplot_mosaic(self, layout, *, subplot_kw=None, gridspec_kw=None,
-                       empty_sentinel='.'):
+    def subplot_mosaic(self, layout, *, sharex=False, sharey=False,
+                       subplot_kw=None, gridspec_kw=None, empty_sentinel='.'):
         """
         Build a layout of Axes based on ASCII art or nested lists.
 
@@ -1684,7 +1683,6 @@ def subplot_mosaic(self, layout, *, subplot_kw=None, gridspec_kw=None,
            This API is provisional and may be revised in the future based on
            early user feedback.
 
-
         Parameters
         ----------
         layout : list of list of {hashable or nested} or str
@@ -1695,7 +1693,7 @@ def subplot_mosaic(self, layout, *, subplot_kw=None, gridspec_kw=None,
                x = [['A panel', 'A panel', 'edge'],
                     ['C panel', '.',       'edge']]
 
-            Produces 4 Axes:
+            produces 4 Axes:
 
             - 'A panel' which is 1 row high and spans the first two columns
             - 'edge' which is 2 rows high and is on the right edge
@@ -1721,6 +1719,12 @@ def subplot_mosaic(self, layout, *, subplot_kw=None, gridspec_kw=None,
             The string notation allows only single character Axes labels and
             does not support nesting but is very terse.
 
+        sharex, sharey : bool, default: False
+            If True, the x-axis (*sharex*) or y-axis (*sharey*) will be shared
+            among all subplots.  In that case, tick label visibility and axis
+            units behave as for `subplots`.  If False, each subplot's x- or
+            y-axis will be independent.
+
         subplot_kw : dict, optional
             Dictionary with keywords passed to the `.Figure.add_subplot` call
             used to create each subplot.
@@ -1748,6 +1752,8 @@ def subplot_mosaic(self, layout, *, subplot_kw=None, gridspec_kw=None,
         # special-case string input
         if isinstance(layout, str):
             layout = self._normalize_grid_string(layout)
+        # Only accept strict bools to allow a possible future API expansion.
+        _api.check_isinstance(bool, sharex=sharex, sharey=sharey)
 
         def _make_array(inp):
             """
@@ -1905,6 +1911,14 @@ def _do_layout(gs, layout, unique_ids, nested):
         rows, cols = layout.shape
         gs = self.add_gridspec(rows, cols, **gridspec_kw)
         ret = _do_layout(gs, layout, *_identify_keys_and_nested(layout))
+        ax0 = next(iter(ret.values()))
+        for ax in ret.values():
+            if sharex:
+                ax.sharex(ax0)
+                ax._label_outer_xaxis()
+            if sharey:
+                ax.sharey(ax0)
+                ax._label_outer_yaxis()
         for k, ax in ret.items():
             if isinstance(k, str):
                 ax.set_label(k)
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -1432,8 +1432,9 @@ def subplots(nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True,
     return fig, axs
 
 
-def subplot_mosaic(layout, *, subplot_kw=None, gridspec_kw=None,
-                   empty_sentinel='.', **fig_kw):
+def subplot_mosaic(layout, *, sharex=False, sharey=False,
+                   subplot_kw=None, gridspec_kw=None, empty_sentinel='.',
+                   **fig_kw):
     """
     Build a layout of Axes based on ASCII art or nested lists.
 
@@ -1444,7 +1445,6 @@ def subplot_mosaic(layout, *, subplot_kw=None, gridspec_kw=None,
        This API is provisional and may be revised in the future based on
        early user feedback.
 
-
     Parameters
     ----------
     layout : list of list of {hashable or nested} or str
@@ -1455,7 +1455,7 @@ def subplot_mosaic(layout, *, subplot_kw=None, gridspec_kw=None,
            x = [['A panel', 'A panel', 'edge'],
                 ['C panel', '.',       'edge']]
 
-        Produces 4 axes:
+        produces 4 axes:
 
         - 'A panel' which is 1 row high and spans the first two columns
         - 'edge' which is 2 rows high and is on the right edge
@@ -1476,6 +1476,12 @@ def subplot_mosaic(layout, *, subplot_kw=None, gridspec_kw=None,
         This only allows only single character Axes labels and does
         not allow nesting but is very terse.
 
+    sharex, sharey : bool, default: False
+        If True, the x-axis (*sharex*) or y-axis (*sharey*) will be shared
+        among all subplots.  In that case, tick label visibility and axis units
+        behave as for `subplots`.  If False, each subplot's x- or y-axis will
+        be independent.
+
     subplot_kw : dict, optional
         Dictionary with keywords passed to the `.Figure.add_subplot` call
         used to create each subplot.
@@ -1507,9 +1513,8 @@ def subplot_mosaic(layout, *, subplot_kw=None, gridspec_kw=None,
     """
     fig = figure(**fig_kw)
     ax_dict = fig.subplot_mosaic(
-        layout,
-        subplot_kw=subplot_kw,
-        gridspec_kw=gridspec_kw,
+        layout, sharex=sharex, sharey=sharey,
+        subplot_kw=subplot_kw, gridspec_kw=gridspec_kw,
         empty_sentinel=empty_sentinel
     )
     return fig, ax_dict
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -883,6 +883,20 @@ def test_nested_user_order(self):
         assert list(ax_dict) == list("ABCDEFGHI")
         assert list(fig.axes) == list(ax_dict.values())
 
+    def test_share_all(self):
+        layout = [
+            ["A", [["B", "C"],
+                   ["D", "E"]]],
+            ["F", "G"],
+            [".", [["H", [["I"],
+                          ["."]]]]]
+        ]
+        fig = plt.figure()
+        ax_dict = fig.subplot_mosaic(layout, sharex=True, sharey=True)
+        ax_dict["A"].set(xscale="log", yscale="logit")
+        assert all(ax.get_xscale() == "log" and ax.get_yscale() == "logit"
+                   for ax in ax_dict.values())
+
 
 def test_reused_gridspec():
     """Test that these all use the same gridspec"""
