refactored grid_spec.tight_layout and fixed matplotlib#1055

Author: leejjoon

lib/matplotlib/gridspec.py

@@ -22,6 +22,7 @@
 import matplotlib.transforms as mtransforms
 
 import numpy as np
+import warnings
 
 class GridSpecBase(object):
     """
@@ -241,7 +242,8 @@ def update(self, **kwargs):
                             ax._sharey.update_params()
                             ax.set_position(ax._sharey.figbox)
                 else:
-                    if ax.get_subplotspec().get_gridspec() == self:
+                    ss = ax.get_subplotspec().get_topmost_subplotspec()
+                    if ss.get_gridspec() == self:
                         ax.update_params()
                         ax.set_position(ax.figbox)
 
@@ -287,67 +289,24 @@ def tight_layout(self, fig, renderer=None, pad=1.08, h_pad=None, w_pad=None, rec
             labels) will fit into. Default is (0, 0, 1, 1).
         """
 
-        from tight_layout import auto_adjust_subplotpars, get_renderer
+        from tight_layout import (get_subplotspec_list,
+                                  get_tight_layout_figure,
+                                  get_renderer)
+
+        subplotspec_list = get_subplotspec_list(fig.axes, grid_spec=self)
+        if None in subplotspec_list:
+            warnings.warn("This figure includes Axes that are not "
+                          "compatible with tight_layout, so its "
+                          "results might be incorrect.")
 
         if renderer is None:
             renderer = get_renderer(fig)
 
-        subplot_list = []
-        num1num2_list = []
-        subplot_dict = {}
-
-        for ax in fig.axes:
-            locator = ax.get_axes_locator()
-            if hasattr(locator, "get_subplotspec"):
-                subplotspec = locator.get_subplotspec()
-            elif hasattr(ax, "get_subplotspec"):
-                subplotspec = ax.get_subplotspec()
-            else:
-                continue
-
-            if subplotspec.get_gridspec() != self: continue
-
-            subplots = subplot_dict.get(subplotspec, [])
-
-            if not subplots:
-                _, _, num1, num2 = subplotspec.get_geometry()
-                num1num2_list.append((num1, num2))
-                subplot_list.append(subplots)
-
-            subplots.append(ax)
-
-
-        kwargs = auto_adjust_subplotpars(fig, renderer,
-                                          nrows_ncols=self.get_geometry(),
-                                          num1num2_list=num1num2_list,
-                                          subplot_list=subplot_list,
-                                          pad=pad, h_pad=h_pad, w_pad=w_pad)
-
-        if rect is not None:
-            # if rect is given, the whole subplots area (including
-            # labels) will fit into the rect instead of the
-            # figure. Note that the rect argument of
-            # *auto_adjust_subplotpars* specify the area that will be
-            # covered by the total area of axes.bbox. Thus we call
-            # auto_adjust_subplotpars twice, where the second run
-            # with adjusted rect parameters.
-
-            left, bottom, right, top = rect
-            if left is not None: left += kwargs["left"]
-            if bottom is not None: bottom += kwargs["bottom"]
-            if right is not None: right -= (1 - kwargs["right"])
-            if top is not None: top -= (1 - kwargs["top"])
-
-            #if h_pad is None: h_pad = pad
-            #if w_pad is None: w_pad = pad
-
-            kwargs = auto_adjust_subplotpars(fig, renderer,
-                                             nrows_ncols=self.get_geometry(),
-                                             num1num2_list=num1num2_list,
-                                             subplot_list=subplot_list,
-                                             pad=pad, h_pad=h_pad, w_pad=w_pad,
-                                             rect=(left, bottom, right, top))
-
+        kwargs = get_tight_layout_figure(fig, fig.axes, subplotspec_list,
+                                         renderer,
+                                         pad=pad, h_pad=h_pad, w_pad=w_pad,
+                                         rect=rect,
+                                         )
 
         self.update(**kwargs)
 

lib/matplotlib/tight_layout.py

@@ -209,12 +209,15 @@ def get_renderer(fig):
     return renderer
 
 
-def get_subplotspec_list(axes_list):
+def get_subplotspec_list(axes_list, grid_spec=None):
     """
     Return a list of subplotspec from the given list of axes.  For an
     instance of axes that does not support subplotspec, None is
     inserted in the list.
 
+    If grid_spec is given, None is inserted for those not from
+    the given grid_spec.
+
     """
     subplotspec_list = []
     for ax in axes_list:
@@ -225,11 +228,16 @@ def get_subplotspec_list(axes_list):
         if hasattr(axes_or_locator, "get_subplotspec"):
             subplotspec = axes_or_locator.get_subplotspec()
             subplotspec = subplotspec.get_topmost_subplotspec()
-            if subplotspec.get_gridspec().locally_modified_subplot_params():
+            gs = subplotspec.get_gridspec()
+            if grid_spec is not None:
+                if gs != grid_spec:
+                    subplotspec = None
+            elif gs.locally_modified_subplot_params():
                 subplotspec = None
         else:
             subplotspec = None
 
+
         subplotspec_list.append(subplotspec)
 
     return subplotspec_list