ENH: implement get_tightbbox on Axis3D

tacaswell · tacaswell · commit f18bd14ca8e4 · 2020-05-27T12:28:17.000-04:00
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -279,8 +279,8 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
                        ratio
             'auto'     automatic; fill the position rectangle with data.
                        This will let the ratio between the size of the
-                       (pseudo-) bounding box be free and will not adjust
-                       anything at draw time.
+                       (pseudo-) bounding box be free and will not
+                       adjust anything at draw time.
             ========   =================================================
 
         adjustable : None or {'box', 'datalim'}, optional
@@ -2861,6 +2861,27 @@ def permutation_matrices(n):
 
         return polygons
 
+    def get_tightbbox(self, renderer, call_axes_locator=True,
+                      bbox_extra_artists=None, *, for_layout_only=False):
+        ret = super().get_tightbbox(renderer,
+                                    call_axes_locator=call_axes_locator,
+                                    bbox_extra_artists=bbox_extra_artists,
+                                    for_layout_only=for_layout_only)
+        batch = [ret]
+        if self._axis3don:
+            for axis in self._get_axis_list():
+                if axis.get_visible():
+                    try:
+                        axis_bb = axis.get_tightbbox(
+                            renderer,
+                            for_layout_only=for_layout_only
+                        )
+                    except TypeError:
+                        # in case downstream library has redefined axis:
+                        axis_bb = axis.get_tightbbox(renderer)
+                if axis_bb:
+                    batch.append(axis_bb)
+        return mtransforms.Bbox.union(batch)
 
 docstring.interpd.update(Axes3D=artist.kwdoc(Axes3D))
 docstring.dedent_interpd(Axes3D.__init__)
diff --git a/lib/mpl_toolkits/mplot3d/axis3d.py b/lib/mpl_toolkits/mplot3d/axis3d.py
@@ -3,7 +3,7 @@
 # Parts rewritten by Reinier Heeres <reinier@heeres.eu>
 
 import numpy as np
-
+import matplotlib.transforms as mtransforms
 from matplotlib import (
     artist, lines as mlines, axis as maxis, patches as mpatches, rcParams)
 from . import art3d, proj3d
@@ -398,12 +398,30 @@ def draw(self, renderer):
         renderer.close_group('axis3d')
         self.stale = False
 
-    # TODO: Get this to work properly when mplot3d supports
-    #       the transforms framework.
-    def get_tightbbox(self, renderer):
-        # Currently returns None so that Axis.get_tightbbox
-        # doesn't return junk info.
-        return None
+    # TODO: Get this to work (more) properly when mplot3d supports the
+    #       transforms framework.
+    def get_tightbbox(self, renderer, *, for_layout_only=False):
+        # inherited docstring
+        if not self.get_visible():
+            return
+        # We have to directly access the internal data structures
+        # (and hope they are up to date) because at draw time we
+        # shift the ticks and their labels around in (x, y) space
+        # based on the projection, the current view port, and their
+        # position in 3D space.  If we extend the transforms framework
+        # into 3D we would not need to do this different book keeping
+        # than we do in the normal axis
+        ticks = [*self.majorTicks, *self.minorTicks]
+        bb_1, bb_2 = self._get_tick_bboxes(ticks, renderer)
+        other = []
+
+        if self.line.get_visible():
+            other.append(self.line.get_window_extent(renderer))
+        if (self.label.get_visible() and not for_layout_only and
+                self.label.get_text()):
+            other.append(self.label.get_window_extent(renderer))
+
+        return mtransforms.Bbox.union([*bb_1, *bb_2, *other])
 
     @property
     def d_interval(self):
