Popuplate initial plotly values selections

Having not found a more elegant solution,
we iterate over the figure's data to create the initial values in the marimo wrapper side.

Author: yairchu

marimo/_plugins/ui/_impl/plotly.py

@@ -142,11 +142,45 @@ def add_selection(selection: go.layout.Selection) -> None:
                 hasattr(selection, k) for k in ["x0", "x1", "y0", "y1"]
             ):
                 return
+
             initial_value["range"] = {
                 "x": [selection.x0, selection.x1],
                 "y": [selection.y0, selection.y1],
             }
 
+            # Find points within the selection range
+            selected_points = []
+            selected_indices = []
+
+            x_axes: list[go.layout.XAxis] = []
+            figure.for_each_xaxis(x_axes.append)
+            [x_axis] = x_axes if len(x_axes) == 1 else [None]
+            y_axes: list[go.layout.YAxis] = []
+            figure.for_each_yaxis(y_axes.append)
+            [y_axis] = y_axes if len(y_axes) == 1 else [None]
+
+            for trace in figure.data:
+                x_data = getattr(trace, "x", None)
+                y_data = getattr(trace, "y", None)
+                if x_data is None or y_data is None:
+                    continue
+                for point_idx, (x, y) in enumerate(zip(x_data, y_data)):
+                    if (
+                        selection.x0 <= x <= selection.x1
+                        and selection.y0 <= y <= selection.y1
+                    ):
+                        selected_points.append(
+                            {
+                                axis.title.text: val
+                                for axis, val in [(x_axis, x), (y_axis, y)]
+                                if axis and axis.title.text
+                            }
+                        )
+                        selected_indices.append(point_idx)
+
+            initial_value["points"] = selected_points
+            initial_value["indices"] = selected_indices
+
         figure.for_each_selection(add_selection)
 
         super().__init__(