Exclude totally empty subplots

packages/python/plotly/plotly/basedatatypes.py

@@ -1302,7 +1302,7 @@ def _add_annotation_like(
             # if exclude_empty_subplots is True, check to see if subplot is
             # empty and return if it is
             if exclude_empty_subplots and (
-                not self._subplot_contains_trace(xref, yref)
+                not self._subplot_not_empty(xref, yref, selector=exclude_empty_subplots)
             ):
                 return self
             # in case the user specified they wanted an axis to refer to the
@@ -1993,8 +1993,8 @@ def add_traces(
         if exclude_empty_subplots:
             data = list(
                 filter(
-                    lambda trace: self._subplot_contains_trace(
-                        trace["xaxis"], trace["yaxis"]
+                    lambda trace: self._subplot_not_empty(
+                        trace["xaxis"], trace["yaxis"], exclude_empty_subplots
                     ),
                     data,
                 )
@@ -3873,19 +3873,56 @@ def _has_subplots(self):
         single plot and so this returns False. """
         return self._grid_ref is not None
 
-    def _subplot_contains_trace(self, xref, yref):
-        return any(
-            t == (xref, yref)
-            for t in [
-                # if a trace exists but has no xaxis or yaxis keys, then it
-                # is plotted with xaxis 'x' and yaxis 'y'
-                (
-                    "x" if d["xaxis"] is None else d["xaxis"],
-                    "y" if d["yaxis"] is None else d["yaxis"],
+    def _subplot_not_empty(self, xref, yref, selector="all"):
+        """
+        xref: string representing the axis. Objects in the plot will be checked
+              for this xref (for layout objects) or xaxis (for traces) to
+              determine if they lie in a certain subplot.
+        yref: string representing the axis. Objects in the plot will be checked
+              for this yref (for layout objects) or yaxis (for traces) to
+              determine if they lie in a certain subplot.
+        selector: can be "all" or an iterable containing some combination of
+                  "traces", "shapes", "annotations", "images". Only the presence
+                  of objects specified in selector will be checked. So if
+                  ["traces","shapes"] is passed then a plot we be considered
+                  non-empty if it contains traces or shapes. If
+                  bool(selector) returns False, no checking is performed and
+                  this function returns True. If selector is True, it is
+                  converted to "all".
+        """
+        if not selector:
+            # If nothing to select was specified then a subplot is always deemed non-empty
+            return True
+        if selector == True:
+            selector = "all"
+        if selector == "all":
+            selector = ["traces", "shapes", "annotations", "images"]
+        ret = False
+        for s in selector:
+            if s == "traces":
+                obj = self.data
+                xaxiskw = "xaxis"
+                yaxiskw = "yaxis"
+            elif s in ["shapes", "annotations", "images"]:
+                obj = self.layout[s]
+                xaxiskw = "xref"
+                yaxiskw = "yref"
+            else:
+                obj = None
+            if obj:
+                ret |= any(
+                    t == (xref, yref)
+                    for t in [
+                        # if a object exists but has no xaxis or yaxis keys, then it
+                        # is plotted with xaxis/xref 'x' and yaxis/yref 'y'
+                        (
+                            "x" if d[xaxiskw] is None else d[xaxiskw],
+                            "y" if d[yaxiskw] is None else d[yaxiskw],
+                        )
+                        for d in obj
+                    ]
                 )
-                for d in self.data
-            ]
-        )
+        return ret
 
 
 class BasePlotlyType(object):

packages/python/plotly/plotly/tests/test_core/test_subplots/test_find_nonempty_subplots.py

@@ -0,0 +1,61 @@
+import pytest
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+from itertools import combinations, product
+from functools import reduce
+
+
+def all_combos(it):
+    return list(
+        reduce(
+            lambda a, b: a + b,
+            [list(combinations(it, r)) for r in range(1, len(it))],
+            [],
+        )
+    )
+
+
+def translate_layout_keys(t):
+    xr, yr = t
+    xr = xr.replace("axis", "")
+    yr = yr.replace("axis", "")
+    return (xr, yr)
+
+
+def get_non_empty_subplots(fig, selector):
+    gr = fig._validate_get_grid_ref()
+    nrows = len(gr)
+    ncols = len(gr[0])
+    sp_addresses = product(range(nrows), range(ncols))
+    # assign a number similar to plotly's xref/yref (e.g, xref=x2) to each
+    # subplot address (xref=x -> 1, but xref=x3 -> 3)
+    # sp_ax_numbers=range(1,len(sp_addresses)+1)
+    # Get those subplot numbers which contain something
+    ret = list(
+        filter(
+            lambda sp: fig._subplot_not_empty(
+                *translate_layout_keys(sp.layout_keys), selector=selector
+            ),
+            [gr[r][c][0] for r, c in sp_addresses],
+        )
+    )
+    return ret
+
+
+def test_choose_correct_non_empty_subplots():
+    # This checks to see that the correct subplots are selected for all
+    # combinations of selectors
+    fig = make_subplots(2, 2)
+    fig.add_trace(go.Scatter(x=[1, 2], y=[3, 4]), row=1, col=1)
+    fig.add_shape(dict(type="rect", x0=1, x1=2, y0=3, y1=4), row=1, col=2)
+    fig.add_annotation(dict(text="A", x=1, y=2), row=2, col=1)
+    fig.add_layout_image(
+        dict(source="test", x=1, y=2, sizex=0.5, sizey=0.5), row=2, col=2
+    )
+    all_subplots = get_non_empty_subplots(fig, "all")
+    selectors = all_combos(["traces", "shapes", "annotations", "images"])
+    subplot_combos = all_combos(all_subplots)
+    assert len(selectors) == len(subplot_combos)
+    for s, spc in zip(selectors, subplot_combos):
+        sps = tuple(get_non_empty_subplots(fig, s))
+        assert sps == spc