finish code and add tests

aswarner · aswarner · commit 8b27ae368985 · 2022-04-04T15:04:24.000-04:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -3004,24 +3004,21 @@ def _update_title_position(self, renderer):
                         (0., 2 * top - title.get_window_extent(renderer).ymin))
                     title.set_position((x, y))
 
+        grouped_axs = self.figure._align_title_groups.get_siblings(self)
         ymax = max(title.get_position()[1] for title in titles)
         for title in titles:
             # now line up all the titles at the highest baseline.
             x, _ = title.get_position()
             title.set_position((x, ymax))
 
-        # Assign
-        grouped_axs = self.figure._align_title_groups.get_siblings(self)
-        ymax = None
+        # Align bboxes of grouped axes to highest in group
+        bb_ymax = None
         ax_max = None
         for ax in grouped_axs:
-            if ymax is None or ax.bbox.ymax > ymax:
-                ymax = ax.bbox.ymax
+            if bb_ymax is None or ax.bbox.ymax > bb_ymax:
+                bb_ymax = ax.bbox.ymax
                 ax_max = ax
-        # print("SELF FIRST: ", self.bbox.xmax)
-        # print("OTHER: ", ax_max.bbox.xmax)
         self.bbox = ax_max.bbox
-        # print("SELF AFTER: ", self.bbox.xmax)
 
     # Drawing
     @martist.allow_rasterization
diff --git a/lib/matplotlib/tests/test_figure.py b/lib/matplotlib/tests/test_figure.py
@@ -100,6 +100,39 @@ def test_align_labels_stray_axes():
     np.testing.assert_allclose(yn[::2], yn[1::2])
 
 
+# TODO add image comparison
+@image_comparison(['figure_align_titles'], extensions=['png', 'svg'],
+                  tol=0 if platform.machine() == 'x86_64' else 0.01)
+def test_align_titles():
+    fig, axs = plt.subplots(2, 2,
+                            subplot_kw={"xlabel": "x", "ylabel": "",
+                                        "title": "Title"})
+    axs[0][0].imshow(plt.np.zeros((5, 3)))
+    axs[0][1].imshow(plt.np.zeros((3, 5)))
+    axs[1][0].imshow(plt.np.zeros((2, 1)))
+    axs[1][1].imshow(plt.np.zeros((1, 2)))
+
+    axs[0][1].set_title('Title2', loc="left")
+    axs[0][1].set_title()
+
+    fig.align_titles()
+
+
+## TODO add image comparison
+@image_comparison(['figure_align_titles_some'], extensions=['png', 'svg'],
+                  tol=0 if platform.machine() == 'x86_64' else 0.01)
+def test_align_titles_param():
+    fig, axs = plt.subplots(2, 2,
+                            subplot_kw={"xlabel": "x", "ylabel": "",
+                                        "title": "t"})
+    axs[0][0].imshow(plt.np.zeros((3, 5)))
+    axs[0][1].imshow(plt.np.zeros((5, 3)))
+    axs[1][0].imshow(plt.np.zeros((2, 1)))
+    axs[1][1].imshow(plt.np.zeros((1, 2)))
+
+    fig.align_titles([axs[0][0], axs[0][1]])
+
+
 def test_figure_label():
     # pyplot figure creation, selection, and closing with label/number/instance
     plt.close('all')
