Fix selection helper iterator

Multiple iterations might happen at the same time, so it's important
that we don't store the iterator state within the selection itself.

Signed-off-by: Bernát Gábor <[email protected]>

Author: gaborbernat

ipydatagrid/datagrid.py

@@ -1,6 +1,7 @@
 # Copyright (c) Bloomberg.
 # Distributed under the terms of the Modified BSD License.
 
+from collections.abc import Iterator
 from copy import deepcopy
 from math import floor
 
@@ -23,26 +24,11 @@
 from .cellrenderer import CellRenderer, TextRenderer
 
 
-class SelectionHelper:
-
-    """A Helper Class for processing selections. Provides an iterator
-    to traverse selected cells.
-    """
-
-    def __init__(self, grid, **kwargs):
-        super().__init__(**kwargs)
-        self._grid = grid
-        self._num_columns = -1
-        self._num_rows = -1
-
-    def __iter__(self):
+class SelectionIterator(Iterator):
+    def __init__(self, selections):
         self._rect_index = 0
         self._cell_index = 0
-        self._selections = [
-            self._transform_rect_for_selection_mode(rect)
-            for rect in self._grid.selections
-        ]
-        return self
+        self._selections = selections
 
     def __next__(self):
         if self._rect_index >= len(self._selections):
@@ -61,6 +47,51 @@ def __next__(self):
         else:
             return row_col
 
+    @staticmethod
+    def _index_to_row_col(rect, index):
+        num_rows = rect["r2"] - rect["r1"] + 1
+        num_cols = rect["c2"] - rect["c1"] + 1
+        if index > (num_rows * num_cols - 1):
+            return None
+
+        return {
+            "r": rect["r1"] + floor(index / num_cols),
+            "c": rect["c1"] + index % num_cols,
+        }
+
+    def _cell_in_previous_selected_rects(self, cell):
+        return any(
+            self._cell_in_rect(cell, self._selections[i])
+            for i in range(0, self._rect_index)
+        )
+
+    @staticmethod
+    def _cell_in_rect(cell, rect):
+        return (
+            rect["r1"] <= cell["r"] <= rect["r2"]
+            and rect["c1"] <= cell["c"] <= rect["c2"]
+        )
+
+
+class SelectionHelper:
+
+    """A Helper Class for processing selections. Provides an iterator
+    to traverse selected cells.
+    """
+
+    def __init__(self, grid, **kwargs):
+        super().__init__(**kwargs)
+        self._grid = grid
+        self._num_columns = -1
+        self._num_rows = -1
+
+    def __iter__(self):
+        selections = [
+            self._transform_rect_for_selection_mode(rect)
+            for rect in self._grid.selections
+        ]
+        return SelectionIterator(selections)
+
     def __len__(self):
         return sum(1 for _ in self)
 
@@ -70,7 +101,7 @@ def all(self):
         represented as a dictionary
         with keys 'r': row and 'c': column
         """
-        return [cell for cell in self]  # noqa: C416
+        return list(self)
 
     def all_values(self):
         """
@@ -81,31 +112,6 @@ def all_values(self):
             for cell in self
         ]
 
-    @staticmethod
-    def _cell_in_rect(cell, rect):
-        return (
-            rect["r1"] <= cell["r"] <= rect["r2"]
-            and rect["c1"] <= cell["c"] <= rect["c2"]
-        )
-
-    def _cell_in_previous_selected_rects(self, cell):
-        return any(
-            self._cell_in_rect(cell, self._selections[i])
-            for i in range(0, self._rect_index)
-        )
-
-    @staticmethod
-    def _index_to_row_col(rect, index):
-        num_rows = rect["r2"] - rect["r1"] + 1
-        num_cols = rect["c2"] - rect["c1"] + 1
-        if index > (num_rows * num_cols - 1):
-            return None
-
-        return {
-            "r": rect["r1"] + floor(index / num_cols),
-            "c": rect["c1"] + index % num_cols,
-        }
-
     def _transform_rect_for_selection_mode(self, rect):
         selection_mode = self._grid.selection_mode
         if selection_mode == "row":

tests/test_datagrid.py

@@ -1,23 +1,22 @@
 import pandas as pd
+import pytest
 
 from ipydatagrid import DataGrid
 
 
-def test_selections():
+@pytest.mark.parametrize("clear", [True, False])
+def test_selections(clear: bool) -> None:
     df = pd.DataFrame(data={"A": [1, 2, 3], "B": [4, 5, 6]})
-    datagrid = DataGrid(
-        df, selection_mode="cell", layout={"height": "100px"}, editable=True
-    )
-    datagrid.select(1, 0, 2, 1)  # Select 1A to 2B
-
-    assert datagrid.selections == [{"r1": 1, "c1": 0, "r2": 2, "c2": 1}]
-
-
-def test_selection_clearing():
-    df = pd.DataFrame(data={"A": [1, 2, 3], "B": [4, 5, 6]})
-    datagrid = DataGrid(
-        df, selection_mode="cell", layout={"height": "100px"}, editable=True
-    )
-    datagrid.select(1, 0, 2, 1)  # Select 1A to 2B
-    datagrid.clear_selection()
-    assert datagrid.selections == []
+    layout = {"height": "100px"}
+    grid = DataGrid(df, selection_mode="cell", layout=layout, editable=True)
+    grid.select(1, 0, 2, 1)  # Select 1A to 2B
+    if clear:
+        grid.clear_selection()
+        assert grid.selected_cells == []
+    else:
+        assert grid.selected_cells == [
+            {"c": 0, "r": 1},
+            {"c": 1, "r": 1},
+            {"c": 0, "r": 2},
+            {"c": 1, "r": 2},
+        ]