Repository: marimo-team/marimo
Pull request: Add indexes property to mo.ui.plotly, and update points to use original column names
Fixes https://github.com/marimo-team/marimo/issues/353

* Add indexes property to mo.ui.plotly - this is a great escape hatch to filter an original dataset in case not all the data is passed to plotly)
* update points to use original column names, instead of `x`, `y`


Problem statement:
mo.ui.plotly: send tooltip back in selection
### Description

One of our users was using reactive plotly.

Their plot had some tooltips, and they would like the tooltips to be sent back with selection.

### Suggested solution

Include a "tooltip" key in the entries of `fig.value`

```
{"x": ..., "y": ..., "tooltip": ...}
```

### Alternative

_No response_

### Additional context

The user chose plotly (and was very happy to see it supported!) because it gives more flexibility than altair — they were plotting a graph with nodes and edges, which altair/vega does not support as far as I can tell.

Hints from issue discussion:

Commit message: Add indexes property to mo.ui.plotly, and update points to use original column names

Patch to review:
diff --git a/frontend/src/plugins/impl/plotly/PlotlyPlugin.tsx b/frontend/src/plugins/impl/plotly/PlotlyPlugin.tsx
index 6b4aa8af753..19d70aaedd8 100644
--- a/frontend/src/plugins/impl/plotly/PlotlyPlugin.tsx
+++ b/frontend/src/plugins/impl/plotly/PlotlyPlugin.tsx
@@ -7,9 +7,9 @@ import Plot, { Figure } from "react-plotly.js";
 import { Logger } from "@/utils/Logger";
 
 import "./plotly.css";
-import { Objects } from "@/utils/objects";
 import { memo, useMemo } from "react";
 import useEvent from "react-use-event-hook";
+import { PlotlyTemplateParser, createParser } from "./parse-from-template";
 
 interface Data {
   figure: Figure;
@@ -21,6 +21,7 @@ type AxisDatum = unknown;
 type T =
   | {
       points?: Array<Record<AxisName, AxisDatum>>;
+      indexes?: number[];
       range?: {
         x?: number[];
         y?: number[];
@@ -79,9 +80,14 @@ export const PlotlyComponent = memo(
         {...figure}
         layout={layout}
         config={config}
-        onSelected={useEvent((evt) => {
+        onSelected={useEvent((evt: Readonly<Plotly.PlotSelectionEvent>) => {
+          if (!evt) {
+            return;
+          }
+
           setValue({
             points: extractPoints(evt.points),
+            indexes: evt.points.map((point) => point.pointIndex),
             range: evt.range,
           });
         })}
@@ -97,7 +103,11 @@ export const PlotlyComponent = memo(
 );
 PlotlyComponent.displayName = "PlotlyComponent";
 
-const keysToInclude = new Set(["x", "y", "value", "label"]);
+/**
+ * This is a hack to extract the points with their original keys,
+ * instead of the ones that Plotly uses internally,
+ * by using the hovertemplate.
+ */
 function extractPoints(
   points: Plotly.PlotDatum[]
 ): Array<Record<AxisName, AxisDatum>> {
@@ -105,14 +115,17 @@ function extractPoints(
     return [];
   }
 
+  let parser: PlotlyTemplateParser | undefined;
+
   return points.map((point) => {
-    const data: Record<AxisName, AxisDatum> = Objects.filter(point, (_v, key) =>
-      keysToInclude.has(key)
-    );
-    // Add name if it exists
-    if (point.data.name) {
-      data.name = point.data.name;
-    }
-    return data;
+    // Get the first hovertemplate
+    const hovertemplate = Array.isArray(point.data.hovertemplate)
+      ? point.data.hovertemplate[0]
+      : point.data.hovertemplate;
+    // Update or create a parser
+    parser = parser
+      ? parser.update(hovertemplate)
+      : createParser(hovertemplate);
+    return parser.parse(point);
   });
 }
diff --git a/frontend/src/plugins/impl/plotly/__tests__/parse-from-template.test.ts b/frontend/src/plugins/impl/plotly/__tests__/parse-from-template.test.ts
new file mode 100644
index 00000000000..1da214bddd9
--- /dev/null
+++ b/frontend/src/plugins/impl/plotly/__tests__/parse-from-template.test.ts
@@ -0,0 +1,73 @@
+/* Copyright 2023 Marimo. All rights reserved. */
+import { describe, expect, it } from "vitest";
+import { createParser } from "../parse-from-template";
+
+describe("createParser", () => {
+  it("does not create a new parser when the template is the same", () => {
+    const parser1 = createParser("template");
+    const parser2 = parser1.update("template");
+    expect(parser1).toBe(parser2);
+
+    const parser3 = parser1.update("template1");
+    expect(parser1).not.toBe(parser3);
+  });
+
+  it("should correctly parse data using the extracted key-selector pairs", () => {
+    const hovertemplate =
+      "Origin=%{customdata[1]}<br>Horsepower=%{x}<br>Miles_per_Gallon=%{y}<br>Weight_in_lbs=%{marker.size}<br>Name=%{customdata[0]}<extra></extra>";
+    const parser = createParser(hovertemplate);
+    const data = {
+      customdata: ["Mustang", "USA"],
+      x: "300",
+      y: "30",
+      "marker.size": "2000",
+    } as unknown as Plotly.PlotDatum;
+    const result = parser.parse(data);
+    expect(result).toEqual({
+      Origin: "USA",
+      Horsepower: "300",
+      Miles_per_Gallon: "30",
+      Weight_in_lbs: "2000",
+      Name: "Mustang",
+    });
+  });
+
+  it("should handle hover templates with no key-selector pairs", () => {
+    const hovertemplate = "No data available";
+    const parser = createParser(hovertemplate);
+    const data = {} as unknown as Plotly.PlotDatum;
+    const result = parser.parse(data);
+    expect(result).toEqual({});
+  });
+
+  it("should handle data with missing keys", () => {
+    const hovertemplate =
+      "Origin=%{customdata[1]}<br>Name=%{customdata[0]}<extra></extra>";
+    const parser = createParser(hovertemplate);
+    const data = {
+      customdata: ["Mustang"], // Missing 'customdata[1]'
+    } as unknown as Plotly.PlotDatum;
+    const result = parser.parse(data);
+    expect(result).toEqual({
+      Origin: undefined, // 'customdata[1]' is missing, so the value is undefined
+      Name: "Mustang",
+    });
+  });
+
+  it("should handle data with nested keys", () => {
+    const hovertemplate =
+      "Model=%{car.model}<br>Year=%{car.year}<extra></extra>";
+    const parser = createParser(hovertemplate);
+    const data = {
+      car: {
+        model: "Mustang",
+        year: "1964",
+      },
+    } as unknown as Plotly.PlotDatum;
+    const result = parser.parse(data);
+    expect(result).toEqual({
+      Model: "Mustang",
+      Year: "1964",
+    });
+  });
+});
diff --git a/frontend/src/plugins/impl/plotly/parse-from-template.ts b/frontend/src/plugins/impl/plotly/parse-from-template.ts
new file mode 100644
index 00000000000..30ffaac4203
--- /dev/null
+++ b/frontend/src/plugins/impl/plotly/parse-from-template.ts
@@ -0,0 +1,55 @@
+/* Copyright 2023 Marimo. All rights reserved. */
+
+import { get } from "lodash-es";
+
+export interface PlotlyTemplateParser {
+  /**
+   * Update the parser with a new hover template, if needed.
+   */
+  update(newTemplate: string): PlotlyTemplateParser;
+  /**
+   * Given a Plotly.PlotDatum, parse the data and return an object with the key-value pairs.
+   */
+  parse(data: Plotly.PlotDatum): Record<string, string>;
+}
+
+/**
+ * Hover template looks something like this:
+ *   "Origin=%{customdata[1]}<br>Horsepower=%{x}<br>Miles_per_Gallon=%{y}<br>Weight_in_lbs=%{marker.size}<br>Name=%{customdata[0]}<extra></extra>"
+ *
+ * We want to parse the keys/selectors to get the values for each data point.
+ *
+ * NOTE: This is a very naive implementation.
+ * It only works for hover templates that have the pattern key=%{selector}.
+ */
+export function createParser(hovertemplate: string): PlotlyTemplateParser {
+  // Regular expression to match the pattern key=%{selector}
+  const regex = /(\w+)=%{([^}]+)}/g;
+
+  // Create an object to hold the key-selector pairs
+  const keySelectorPairs: Record<string, string> = {};
+
+  let match: RegExpExecArray | null;
+
+  // Use exec to find all occurrences of the pattern and iterate over them
+  while ((match = regex.exec(hovertemplate)) !== null) {
+    // match[1] is the key, match[2] is the selector
+    keySelectorPairs[match[1]] = match[2];
+  }
+
+  return {
+    update(newTemplate: string): PlotlyTemplateParser {
+      // If the template hasn't changed, return the same parser
+      if (newTemplate === hovertemplate) {
+        return this;
+      }
+      return createParser(newTemplate);
+    },
+    parse(data: Plotly.PlotDatum) {
+      return Object.entries(keySelectorPairs).reduce((acc, [key, selector]) => {
+        acc[key] = get(data, selector);
+        return acc;
+      }, {} as Record<string, string>);
+    },
+  };
+}
diff --git a/marimo/_plugins/ui/_impl/plotly.py b/marimo/_plugins/ui/_impl/plotly.py
index 0aa91ab96cd..2d544929607 100644
--- a/marimo/_plugins/ui/_impl/plotly.py
+++ b/marimo/_plugins/ui/_impl/plotly.py
@@ -33,6 +33,7 @@
 #     "field1": [min, max],
 #     "field2": [min, max],
 #   },
+#  "indexes": int[],
 # }
 PlotlySelection = Dict[str, JSONType]
 
@@ -120,6 +121,14 @@ def points(self) -> List[Dict[str, Any]]:
             return []
         return self._selection_data["points"]  # type:ignore
 
+    @property
+    def indexes(self) -> List[int]:
+        if not self._selection_data:
+            return []
+        if "indexes" not in self._selection_data:
+            return []
+        return self._selection_data["indexes"]  # type:ignore
+
     def _convert_value(self, value: PlotlySelection) -> Any:
         # Store the selection data
         self._selection_data = value
diff --git a/marimo/_smoke_tests/third_party/plotly.py b/marimo/_smoke_tests/third_party/plotly.py
index 6b4c5ff1fbd..36b67d9e78a 100644
--- a/marimo/_smoke_tests/third_party/plotly.py
+++ b/marimo/_smoke_tests/third_party/plotly.py
@@ -1,7 +1,7 @@
 # Copyright 2023 Marimo. All rights reserved.
 import marimo
 
-__generated_with = "0.1.47"
+__generated_with = "0.1.51"
 app = marimo.App(width="full")
 
 
@@ -35,19 +35,24 @@ def __(mo, px):
 
 @app.cell
 def __(mo, plot):
-    mo.hstack(
+    mo.vstack(
         [
-            mo.ui.table(plot.value, label="Points", selection=None),
-            mo.ui.table(
+            mo.hstack(
                 [
-                    {"start": r[0], "end": r[1], "axis": key}
-                    for key, r in plot.ranges.items()
+                    mo.ui.table(plot.value, label="Points", selection=None),
+                    mo.ui.table(
+                        [
+                            {"start": r[0], "end": r[1], "axis": key}
+                            for key, r in plot.ranges.items()
+                        ],
+                        selection=None,
+                        label="Ranges",
+                    ),
                 ],
-                selection=None,
-                label="Ranges",
+                widths="equal",
             ),
-        ],
-        widths="equal",
+            plot.indexes,
+        ]
     )
     return
 
@@ -144,7 +149,7 @@ def __(cars, mo, px, sample_size):
         y="Miles_per_Gallon",
         color="Origin",
         size="Weight_in_lbs",
-        hover_data=["Name"],
+        hover_data=["Name", "Origin"],
     )
 
     _fig


Code review comments:
diff_hunk: @@ -120,6 +121,14 @@ def points(self) -> List[Dict[str, Any]]:
             return []
         return self._selection_data["points"]  # type:ignore
 
+    @property
+    def indexes(self) -> List[int]:
path: marimo/_plugins/ui/_impl/plotly.py
text: @user1:
nit: let's make this "indices"