For custom-sized subplots, _add_annotation_like works with "all"

The implementation doesn't try to put shapes on subplots at are None in the
grid_ref.

A test is there that goes with it.

Author: nicholas-esterer

packages/python/plotly/plotly/basedatatypes.py

@@ -1136,10 +1136,9 @@ def _add_annotation_like(
 
         # Address multiple subplots
         if row is not None and _is_select_subplot_coordinates_arg(row, col):
-            # TODO add product argument
-            rows, cols = self._select_subplot_coordinates(row, col)
-            # TODO do we have to unzip the row and columns, just to zip them again?
-            for r, c in zip(rows, cols):
+            # TODO product argument could be added
+            rows_cols = self._select_subplot_coordinates(row, col)
+            for r, c in rows_cols:
                 self._add_annotation_like(
                     prop_singular,
                     prop_plural,
@@ -1732,9 +1731,8 @@ def add_trace(self, trace, row=None, col=None, secondary_y=None):
         # Address multiple subplots
         if row is not None and _is_select_subplot_coordinates_arg(row, col):
             # TODO add product argument
-            rows, cols = self._select_subplot_coordinates(row, col)
-            # TODO do we have to unzip the row and columns, just to zip them again?
-            for r, c in zip(rows, cols):
+            rows_cols = self._select_subplot_coordinates(row, col)
+            for r, c in rows_cols:
                 self.add_trace(trace, row=r, col=c, secondary_y=secondary_y)
             return self
 
@@ -1978,8 +1976,10 @@ def _select_subplot_coordinates(self, rows, cols, product=False):
             [rows, cols], list(self._get_subplot_rows_columns()), product=product,
         )
         t = list(t)
-        r, c = _unzip_pairs(t)
-        return (r, c)
+        # remove rows and cols where the subplot is "None"
+        grid_ref = self._validate_get_grid_ref()
+        t = list(filter(lambda u: grid_ref[u[0] - 1][u[1] - 1] is not None, t))
+        return t
 
     def get_subplot(self, row, col, secondary_y=False):
         """

packages/python/plotly/plotly/tests/test_core/test_update_objects/test_paper_span_shapes.py

@@ -383,6 +383,96 @@ def test_non_subplot_add_span_shape(test_input, expected, non_subplot_fig_fixtur
     _check_figure_shapes(test_input, expected, non_subplot_fig_fixture)
 
 
-def test_invalid_subplot_address(subplot_fig_fixture):
+@pytest.mark.parametrize(
+    "test_input",
+    [
+        (go.Figure.add_hline, dict(y=10, row=4, col=5)),
+        # valid row, invalid column
+        (go.Figure.add_hline, dict(y=10, row=1, col=5)),
+    ],
+)
+def test_invalid_subplot_address(test_input, subplot_fig_fixture):
+    f, kwargs = test_input
     with pytest.raises(IndexError):
-        subplot_fig_fixture.add_hline(y=10, row=4, col=5)
+        f(subplot_fig_fixture, **kwargs)
+
+
+def _check_figure_shapes_custom_sized(test_input, expected, fig):
+    # look up domains in fig
+    corrects = []
+    for d, ax in expected:
+        dom = fig["layout"][ax]["domain"]
+        if ax[: len("xaxis")] == "xaxis":
+            d["x0"], d["x1"] = dom
+        elif ax[: len("yaxis")] == "yaxis":
+            d["y0"], d["y1"] = dom
+        else:
+            raise ValueError("bad axis")
+        corrects.append(d)
+    f, kwargs = test_input
+    f(fig, **kwargs)
+    ret = True
+    for s, d in zip(fig.layout.shapes, corrects):
+        ret &= _cmp_partial_dict(s, d)
+    assert ret
+
+
+@pytest.mark.parametrize(
+    "test_input,expected",
+    # test_input: (function,kwargs)
+    # expected: list of dictionaries with key:value pairs we expect in the added shapes
+    [
+        (
+            (go.Figure.add_vline, dict(x=1.5, row="all", col=2)),
+            [
+                (
+                    {
+                        "type": "line",
+                        "x0": 1.5,
+                        "x1": 1.5,
+                        "xref": "x2",
+                        "yref": "paper",
+                    },
+                    "yaxis2",
+                ),
+                (
+                    {
+                        "type": "line",
+                        "x0": 1.5,
+                        "x1": 1.5,
+                        "xref": "x6",
+                        "yref": "paper",
+                    },
+                    "yaxis6",
+                ),
+            ],
+        ),
+        (
+            (go.Figure.add_hline, dict(y=1.5, row=5, col="all")),
+            [
+                (
+                    {
+                        "type": "line",
+                        "yref": "y5",
+                        "y0": 1.5,
+                        "y1": 1.5,
+                        "xref": "paper",
+                    },
+                    "xaxis5",
+                ),
+                (
+                    {
+                        "type": "line",
+                        "yref": "y6",
+                        "y0": 1.5,
+                        "y1": 1.5,
+                        "xref": "paper",
+                    },
+                    "xaxis6",
+                ),
+            ],
+        ),
+    ],
+)
+def test_custom_sized_subplots(test_input, expected, custom_sized_subplots):
+    _check_figure_shapes_custom_sized(test_input, expected, custom_sized_subplots)

packages/python/plotly/plotly/tests/test_core/test_update_objects/test_select_subplots.py

@@ -176,6 +176,7 @@ def _sort_row_col_lists(rows, cols):
 def test_select_subplot_coordinates(subplot_fig_fixture, test_input, expected):
     rows, cols, product = test_input
     er, ec = _sort_row_col_lists(*expected)
-    r, c = subplot_fig_fixture._select_subplot_coordinates(rows, cols, product=product)
+    t = subplot_fig_fixture._select_subplot_coordinates(rows, cols, product=product)
+    r, c = _unzip_pairs(t)
     r, c = _sort_row_col_lists(r, c)
     assert (r == er) and (c == ec)