tight_layout.py: pep8 fix and some refactoring

Author: leejjoon

lib/matplotlib/figure.py

@@ -1422,17 +1422,20 @@ def tight_layout(self, renderer=None, pad=1.08, h_pad=None, w_pad=None, rect=Non
             labels) will fit into. Default is (0, 0, 1, 1).
         """
 
-        from tight_layout import get_renderer, get_tight_layout_figure, \
-             get_subplotspec_list
+        from tight_layout import (get_renderer, get_tight_layout_figure,
+                                  get_subplotspec_list)
 
-        if None in get_subplotspec_list(self.axes):
-            warnings.warn("tight_layout can only process Axes that descend "
-                          "from SubplotBase; results might be incorrect.")
+        subplotspec_list = get_subplotspec_list(self.axes)
+        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(self)
 
-        kwargs = get_tight_layout_figure(self, self.axes, renderer,
+        kwargs = get_tight_layout_figure(self, self.axes, subplotspec_list,
+                                         renderer,
                                          pad=pad, h_pad=h_pad, w_pad=w_pad,
                                          rect=rect)
 

lib/matplotlib/tight_layout.py

@@ -218,25 +218,24 @@ def get_subplotspec_list(axes_list):
     """
     subplotspec_list = []
     for ax in axes_list:
-        locator = ax.get_axes_locator()
-        if hasattr(locator, "get_subplotspec"):
-            subplotspec = locator.get_subplotspec().get_topmost_subplotspec()
-        elif hasattr(ax, "get_subplotspec"):
-            subplotspec = ax.get_subplotspec().get_topmost_subplotspec()
+        axes_or_locator = ax.get_axes_locator()
+        if axes_or_locator is None:
+            axes_or_locator = ax
+
+        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():
+                subplotspec = None
         else:
             subplotspec = None
 
-        if subplotspec is not None and \
-               subplotspec.get_gridspec().locally_modified_subplot_params():
-
-            subplotspec = None
-
         subplotspec_list.append(subplotspec)
 
     return subplotspec_list
 
 
-def get_tight_layout_figure(fig, axes_list, renderer,
+def get_tight_layout_figure(fig, axes_list, subplotspec_list, renderer,
                             pad=1.08, h_pad=None, w_pad=None, rect=None):
     """
     Return subplot parameters for tight-layouted-figure with specified
@@ -248,6 +247,9 @@ def get_tight_layout_figure(fig, axes_list, renderer,
 
       *axes_list* : a list of axes
 
+      *subplotspec_list* : a list of subplotspec associated with each
+        axes in axes_list
+
       *renderer* : renderer instance
 
       *pad* : float
@@ -265,7 +267,6 @@ def get_tight_layout_figure(fig, axes_list, renderer,
     """
 
 
-    subplotspec_list = []
     subplot_list = []
     nrows_list = []
     ncols_list = []
@@ -275,8 +276,10 @@ def get_tight_layout_figure(fig, axes_list, renderer,
                       # same subplot_interface (e.g, axes_grid1). Thus
                       # we need to join them together.
 
+    subplotspec_list2 = []
+
     for ax, subplotspec in zip(axes_list,
-                               get_subplotspec_list(axes_list)):
+                               subplotspec_list):
         if subplotspec is None:
             continue
 
@@ -286,7 +289,7 @@ def get_tight_layout_figure(fig, axes_list, renderer,
             myrows, mycols, _, _ = subplotspec.get_geometry()
             nrows_list.append(myrows)
             ncols_list.append(mycols)
-            subplotspec_list.append(subplotspec)
+            subplotspec_list2.append(subplotspec)
             subplot_list.append(subplots)
             ax_bbox_list.append(subplotspec.get_position(fig))
 
@@ -296,7 +299,7 @@ def get_tight_layout_figure(fig, axes_list, renderer,
     max_ncols = max(ncols_list)
 
     num1num2_list = []
-    for subplotspec in subplotspec_list:
+    for subplotspec in subplotspec_list2:
         rows, cols, num1, num2 = subplotspec.get_geometry()
         div_row, mod_row = divmod(max_nrows, rows)
         div_col, mod_col = divmod(max_ncols, cols)