From 1945b4098a1a4b0b3b4c8bfa7179f65c0d7d835d Mon Sep 17 00:00:00 2001
From: Andrew Stein <steinlink@gmail.com>
Date: Wed, 6 Nov 2024 11:29:53 -0500
Subject: [PATCH] Fix Jupyter widget plugin switch

Signed-off-by: Andrew Stein <steinlink@gmail.com>
---
 .../test/jupyter/widget.spec.mts              | 46 +++++++++++++++++++
 .../src/js/model/toolbar.js                   |  2 +-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/packages/perspective-jupyterlab/test/jupyter/widget.spec.mts b/packages/perspective-jupyterlab/test/jupyter/widget.spec.mts
index cdcffae342..794f8a3681 100644
--- a/packages/perspective-jupyterlab/test/jupyter/widget.spec.mts
+++ b/packages/perspective-jupyterlab/test/jupyter/widget.spec.mts
@@ -442,6 +442,52 @@ perpsective.PerspectiveWidget(df, **config)
             expect(errored).toBe(false);
         });
 
+        test_jupyter(
+            "Toggles to datagrid and back regression",
+            [
+                "w = perspective.widget.PerspectiveWidget(arrow_data, columns=['f64', 'str', 'datetime'])",
+                "w",
+            ],
+            async ({ page }) => {
+                await default_body(page);
+                const num_columns = await page
+                    .locator("regular-table thead tr")
+                    .first()
+                    .evaluate((tr) => tr.childElementCount);
+
+                async function toggle(plugin) {
+                    await page.locator(".plugin-select-item").click();
+                    await page
+                        .locator("#plugin_selector_container.open")
+                        .waitFor();
+
+                    await page
+                        .locator(`[data-plugin=${plugin}].plugin-select-item`)
+                        .click();
+
+                    await page
+                        .locator("#plugin_selector_container:not(.open)")
+                        .waitFor();
+
+                    await page.evaluate(async () => {
+                        await document
+                            .querySelector("perspective-viewer")!
+                            .flush();
+                    });
+                }
+
+                await toggle('"X/Y Line"');
+                await toggle("Datagrid");
+                await toggle('"X/Y Line"');
+                await toggle("Datagrid");
+
+                // expect(num_columns).toEqual(3);
+                await expect(
+                    page.locator("regular-table tbody tr")
+                ).toHaveCount(5);
+            }
+        );
+
         // *************************
         // UTILS
         // *************************
diff --git a/packages/perspective-viewer-datagrid/src/js/model/toolbar.js b/packages/perspective-viewer-datagrid/src/js/model/toolbar.js
index ce6601a900..2177580613 100644
--- a/packages/perspective-viewer-datagrid/src/js/model/toolbar.js
+++ b/packages/perspective-viewer-datagrid/src/js/model/toolbar.js
@@ -27,7 +27,7 @@ export function toggle_edit_mode(mode = undefined) {
     }
 
     this.model._edit_mode = mode;
-    this.parentElement.setSelection();
+    this.parentElement?.setSelection();
     this.model._selection_state = {
         selected_areas: [],
         dirty: true,