Merge branch 'issue-2140-b' into issue-2140

nicholas-esterer · nicholas-esterer · commit 7fca7de1157f · 2020-06-12T15:37:37.000-04:00
diff --git a/.gitignore b/.gitignore
@@ -46,5 +46,4 @@ temp-plot.html
 doc/python/.ipynb_checkpoints
 doc/python/.mapbox_token
 doc/.ipynb_checkpoints
-
 tags
diff --git a/packages/python/plotly/plotly/basedatatypes.py b/packages/python/plotly/plotly/basedatatypes.py
@@ -19,58 +19,41 @@
 Undefined = object()
 
 
-def _rcindex_type(d):
-    all_flag = False
-    if type(d) == type(tuple()):
-        d, f = d
-        if f == "all":
-            all_flag = True
-    if type(d) == type(range(1)):
-        d = list(d)
-    if type(d) == type(int()):
-        return (d, "i", all_flag)
-    elif type(d) == type(list()):
-        return (d, "l", all_flag)
-    elif d == "all":
-        return (d, "a", all_flag)
-    else:
-        raise TypeError(
-            "argument must be 'all', int or list, got {d_type}".format(
-                d_type=str(type(d))
-            )
-        )
+def _unzip_pairs(pairs):
+    pairs = list(pairs)
+    return ([t[0] for t in pairs], [t[1] for t in pairs])
 
 
-def _rcsingle_index_to_list(d):
-    if type(d) == type(int()):
-        return [d]
-    return d
-
-
-def _row_col_index_combinations(rows, cols, max_n_rows, max_n_cols):
-    all_flag = False
-    rows, rtype, f = _rcindex_type(rows)
-    all_flag |= f
-    cols, ctype, f = _rcindex_type(cols)
-    all_flag |= f
-    rows = _rcsingle_index_to_list(rows)
-    cols = _rcsingle_index_to_list(cols)
-    ptype = (rtype, ctype)
-    all_rows = range(1, max_n_rows + 1)
-    all_cols = range(1, max_n_cols + 1)
-    if ptype == ("a", "a"):
-        return list(itertools.product(all_rows, all_cols))
-    elif ptype == ("l", "a") or ptype == ("i", "a"):
-        return list(itertools.product(rows, all_cols))
-    elif ptype == ("a", "l") or ptype == ("a", "i"):
-        return list(itertools.product(all_rows, cols))
-    elif ptype == ("l", "l"):
-        if len(rows) == len(cols) and not all_flag:
-            return list(zip(rows, cols))
-        else:
-            return list(itertools.product(rows, cols))
-    elif ptype == ("l", "i") or ptype == ("i", "i") or ptype == ("i", "l"):
-        return list(itertools.product(rows, cols))
+def _indexing_combinations(dims, alls, product=False):
+    """
+    Gives indexing tuples specified by the coordinates in dims.
+    If a member of dims is 'all' then it is replaced by the corresponding member
+    in alls.
+    If product is True, then the cartesian product of all the indices is
+    returned, otherwise the zip (that means index lists of mis-matched length
+    will yield a list of tuples whose length is the length of the shortest
+    list).
+    """
+    if len(dims) == 0:
+        # this is because list(itertools.product(*[])) returns [()] which has non-zero
+        # length!
+        return []
+    if len(dims) != len(alls):
+        raise ValueError(
+            "Must have corresponding values in alls for each value of dims. Got dims=%s and alls=%s."
+            % (str(dims), str(alls))
+        )
+    r = []
+    for d, a in zip(dims, alls):
+        if d == "all":
+            d = a
+        elif type(d) != type(list()):
+            d = [d]
+        r.append(d)
+    if product:
+        return itertools.product(*r)
+    else:
+        return zip(*r)
 
 
 class BaseFigure(object):
@@ -1918,6 +1901,41 @@ def _validate_get_grid_ref(self):
             )
         return grid_ref
 
+    def _get_subplot_rows_columns(self):
+        """
+        Returns a pair of lists, the first containing all the row indices and
+        the second all the column indices.
+        """
+        # currently, this just iterates over all the rows and columns (because
+        # self._grid_ref is currently always rectangular)
+        grid_ref = self._validate_get_grid_ref()
+        nrows = len(grid_ref)
+        ncols = len(grid_ref[0])
+        return (range(1, nrows + 1), range(1, ncols + 1))
+
+    def _get_subplot_coordinates(self):
+        """
+        Returns an iterator over (row,col) pairs representing all the possible
+        subplot coordinates.
+        """
+        return itertools.product(*self._get_subplot_rows_columns())
+
+    def _select_subplot_coordinates(self, rows, cols, product=False):
+        """
+        Allows selecting all or a subset of the subplots.
+        If any of rows or columns is 'all', product is set to True. This is
+        probably the expected behaviour, so that rows=1,cols='all' selects all
+        the columns in row 1 (otherwise it would just select the subplot in the
+        first row and first column).
+        """
+        product |= any([s == "all" for s in [rows, cols]])
+        t = _indexing_combinations(
+            [rows, cols], list(self._get_subplot_rows_columns()), product=product,
+        )
+        t = list(t)
+        r, c = _unzip_pairs(t)
+        return (r, c)
+
     def get_subplot(self, row, col, secondary_y=False):
         """
         Return an object representing the subplot at the specified row
diff --git a/packages/python/plotly/plotly/tests/test_core/test_update_objects/test_select_subplots.py b/packages/python/plotly/plotly/tests/test_core/test_update_objects/test_select_subplots.py
@@ -0,0 +1,181 @@
+import plotly.graph_objs as go
+from plotly.subplots import make_subplots
+from plotly.basedatatypes import _indexing_combinations, _unzip_pairs
+import pytest
+
+NROWS = 4
+NCOLS = 5
+
+
+@pytest.fixture
+def subplot_fig_fixture():
+    fig = make_subplots(NROWS, NCOLS)
+    return fig
+
+
+@pytest.fixture
+def non_subplot_fig_fixture():
+    fig = go.Figure(go.Scatter(x=[1, 2, 3], y=[4, 3, 2]))
+    return fig
+
+
+def test_invalid_validate_get_grid_ref(non_subplot_fig_fixture):
+    with pytest.raises(Exception):
+        _ = non_subplot_fig_fixture._validate_get_grid_ref()
+
+
+def test_get_subplot_coordinates(subplot_fig_fixture):
+    assert set(subplot_fig_fixture._get_subplot_coordinates()) == set(
+        [(r, c) for r in range(1, NROWS + 1) for c in range(1, NCOLS + 1)]
+    )
+
+
+def test_indexing_combinations_edge_cases():
+    # Although in theory _indexing_combinations works for any number of
+    # dimensions, we're just interested in 2D for subplots so that's what we
+    # test here.
+    assert _indexing_combinations([], []) == []
+    with pytest.raises(ValueError):
+        _ = _indexing_combinations([[1, 2], [3, 4, 5]], [[1, 2]])
+
+
+# 18 combinations of input possible:
+# ('all', 'all', 'product=True'),
+# ('all', 'all', 'product=False'),
+# ('all', '<list>', 'product=True'),
+# ('all', '<list>', 'product=False'),
+# ('all', '<not-list>', 'product=True'),
+# ('all', '<not-list>', 'product=False'),
+# ('<list>', 'all', 'product=True'),
+# ('<list>', 'all', 'product=False'),
+# ('<list>', '<list>', 'product=True'),
+# ('<list>', '<list>', 'product=False'),
+# ('<list>', '<not-list>', 'product=True'),
+# ('<list>', '<not-list>', 'product=False'),
+# ('<not-list>', 'all', 'product=True'),
+# ('<not-list>', 'all', 'product=False'),
+# ('<not-list>', '<list>', 'product=True'),
+# ('<not-list>', '<list>', 'product=False'),
+# ('<not-list>', '<not-list>', 'product=True'),
+# ('<not-list>', '<not-list>', 'product=False')
+# For <not-list> we choose int because that's what the subplot indexing routines
+# will work with.
+all_rows = [1, 2, 3, 4]
+all_cols = [1, 2, 3, 4, 5]
+
+
+@pytest.mark.parametrize(
+    "test_input,expected",
+    [
+        (
+            dict(dims=["all", "all"], alls=[all_rows, all_cols], product=False),
+            set(zip(all_rows, all_cols)),
+        ),
+        (
+            dict(dims=["all", "all"], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in all_rows for c in all_cols]),
+        ),
+        (
+            dict(dims=["all", [2, 4, 5]], alls=[all_rows, all_cols], product=False),
+            set(zip(all_rows, [2, 4, 5])),
+        ),
+        (
+            dict(dims=["all", [2, 4, 5]], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in all_rows for c in [2, 4, 5]]),
+        ),
+        (
+            dict(dims=["all", 3], alls=[all_rows, all_cols], product=False),
+            set([(all_rows[0], 3)]),
+        ),
+        (
+            dict(dims=["all", 3], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in all_rows for c in [3]]),
+        ),
+        (
+            dict(dims=[[1, 3], "all"], alls=[all_rows, all_cols], product=False),
+            set(zip([1, 3], all_cols)),
+        ),
+        (
+            dict(dims=[[1, 3], "all"], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in [1, 3] for c in all_cols]),
+        ),
+        (
+            dict(dims=[[1, 3], [2, 4, 5]], alls=[all_rows, all_cols], product=False),
+            set(zip([1, 3], [2, 4, 5])),
+        ),
+        (
+            dict(dims=[[1, 3], [2, 4, 5]], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in [1, 3] for c in [2, 4, 5]]),
+        ),
+        (
+            dict(dims=[[1, 3], 3], alls=[all_rows, all_cols], product=False),
+            set([(1, 3)]),
+        ),
+        (
+            dict(dims=[[1, 3], 3], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in [1, 3] for c in [3]]),
+        ),
+        (
+            dict(dims=[2, "all"], alls=[all_rows, all_cols], product=False),
+            set([(2, all_cols[0])]),
+        ),
+        (
+            dict(dims=[2, "all"], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in [2] for c in all_cols]),
+        ),
+        (
+            dict(dims=[2, [2, 4, 5]], alls=[all_rows, all_cols], product=False),
+            set([(2, 2)]),
+        ),
+        (
+            dict(dims=[2, [2, 4, 5]], alls=[all_rows, all_cols], product=True),
+            set([(r, c) for r in [2] for c in [2, 4, 5]]),
+        ),
+        (dict(dims=[2, 3], alls=[all_rows, all_cols], product=False), set([(2, 3)])),
+        (dict(dims=[2, 3], alls=[all_rows, all_cols], product=True), set([(2, 3)])),
+    ],
+)
+def test_indexing_combinations(test_input, expected):
+    assert set(_indexing_combinations(**test_input)) == expected
+
+
+def _sort_row_col_lists(rows, cols):
+    # makes sure that row and column lists are compared in the same order
+    # sorted on rows
+    si = sorted(range(len(rows)), key=lambda i: rows[i])
+    rows = [rows[i] for i in si]
+    cols = [cols[i] for i in si]
+    return (rows, cols)
+
+
+# _indexing_combinations tests most cases of the following function
+# we just need to test that setting rows or cols to 'all' makes product True,
+# and if not, we can still set product to True.
+@pytest.mark.parametrize(
+    "test_input,expected",
+    [
+        (
+            ("all", [2, 4, 5], False),
+            _unzip_pairs([(r, c) for r in range(1, NROWS + 1) for c in [2, 4, 5]]),
+        ),
+        (
+            ([1, 3], "all", False),
+            _unzip_pairs([(r, c) for r in [1, 3] for c in range(1, NCOLS + 1)]),
+        ),
+        (
+            ([1, 3], "all", True),
+            _unzip_pairs([(r, c) for r in [1, 3] for c in range(1, NCOLS + 1)]),
+        ),
+        (([1, 3], [2, 4, 5], False), _unzip_pairs([(1, 2), (3, 4)])),
+        (
+            ([1, 3], [2, 4, 5], True),
+            _unzip_pairs([(r, c) for r in [1, 3] for c in [2, 4, 5]]),
+        ),
+    ],
+)
+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)
+    r, c = _sort_row_col_lists(r, c)
+    assert (r == er) and (c == ec)
