From 9c521dddad9cc1ad4eec8dad56b516c31e50e61b Mon Sep 17 00:00:00 2001
From: "Maarten A. Breddels" <maartenbreddels@gmail.com>
Date: Tue, 16 Mar 2021 17:05:28 +0100
Subject: [PATCH] lab support

---
 jupyterlab_widgets/src/manager.ts | 31 +++++++++++++-------
 jupyterlab_widgets/src/plugin.ts  | 47 ++++++++++++++++++++++---------
 2 files changed, 55 insertions(+), 23 deletions(-)

diff --git a/jupyterlab_widgets/src/manager.ts b/jupyterlab_widgets/src/manager.ts
index 73834b23792..0eb672d37ba 100644
--- a/jupyterlab_widgets/src/manager.ts
+++ b/jupyterlab_widgets/src/manager.ts
@@ -517,16 +517,26 @@ export class WidgetManager extends LabWidgetManager {
    * Save the widget state to the context model.
    */
   private _saveState(): void {
-    const visibleWidgets = [];
-    for (const cell of toArray(this._context.model.cells)) {
-      if (cell.type === 'code') {
-        const codeCell = cell as ICodeCellModel;
-        for (let i = 0; i < codeCell.outputs.length; i++) {
-          const output = codeCell.outputs.get(i);
-          if (output.data[WIDGET_VIEW_MIMETYPE]) {
-            const widgetData = output.data[WIDGET_VIEW_MIMETYPE] as any;
-            const modelId = widgetData['model_id'];
-            visibleWidgets.push(modelId);
+    const metadata: any = this._context.model?.metadata.toJSON() || null;
+    // the default is to not store widgets
+    const store = metadata?.widgets?.store || 'none';
+    if (store == 'none') {
+      return;
+    }
+
+    let visibleWidgets = undefined;
+    if (store == 'displayed') {
+      visibleWidgets = [];
+      for (const cell of toArray(this._context.model.cells)) {
+        if (cell.type === 'code') {
+          const codeCell = cell as ICodeCellModel;
+          for (let i = 0; i < codeCell.outputs.length; i++) {
+            const output = codeCell.outputs.get(i);
+            if (output.data[WIDGET_VIEW_MIMETYPE]) {
+              const widgetData = output.data[WIDGET_VIEW_MIMETYPE] as any;
+              const modelId = widgetData['model_id'];
+              visibleWidgets.push(modelId);
+            }
           }
         }
       }
@@ -536,6 +546,7 @@ export class WidgetManager extends LabWidgetManager {
       visibleWidgets: visibleWidgets
     });
     this._context.model.metadata.set('widgets', {
+      store: store,
       'application/vnd.jupyter.widget-state+json': state
     });
   }
diff --git a/jupyterlab_widgets/src/plugin.ts b/jupyterlab_widgets/src/plugin.ts
index cc9274be225..17f9e8d9627 100644
--- a/jupyterlab_widgets/src/plugin.ts
+++ b/jupyterlab_widgets/src/plugin.ts
@@ -247,24 +247,45 @@ function activateWidgetExtension(
     });
   }
 
+  function isNotebookActive(): boolean {
+    return (
+      tracker?.currentWidget !== null &&
+      tracker?.currentWidget === app.shell.currentWidget
+    );
+  }
+
   if (settingRegistry !== null) {
-    // Add a command for automatically saving (jupyter-)widget state.
-    commands.addCommand('@jupyter-widgets/jupyterlab-manager:saveWidgetState', {
-      label: 'Save Widget State Automatically',
-      execute: args => {
-        return settingRegistry
-          .set(plugin.id, 'saveState', !SETTINGS.saveState)
-          .catch((reason: Error) => {
-            console.error(`Failed to set ${plugin.id}: ${reason.message}`);
-          });
-      },
-      isToggled: () => SETTINGS.saveState
-    });
+    // Add commands for changing how we save (jupyter-)widget state.
+    let createStoreCommand = (name: string, value: string, label: string) => {
+      commands.addCommand(`@jupyter-widgets/jupyterlab-manager:${name}`, {
+        label: label,
+        execute: args => {
+          const current = tracker?.currentWidget || null;
+          current?.model?.metadata.set('widgets', { store: value });
+          app.commands.execute('docmanager:save');
+        },
+        isToggled: () => {
+          const current = tracker?.currentWidget || null;
+          const metadata: any = current?.model?.metadata.toJSON() || null;
+          return metadata?.widgets?.store === value;
+        },
+        isEnabled: isNotebookActive
+      });
+    };
+    createStoreCommand('saveNoWidgets', 'none', 'Store no Widgets');
+    createStoreCommand(
+      'saveDisplayedWidgets',
+      'displayed',
+      'Store Displayed Widgets'
+    );
+    createStoreCommand('saveAllWidgets', 'all', 'Store all Widgets');
   }
 
   if (menu) {
     menu.settingsMenu.addGroup([
-      { command: '@jupyter-widgets/jupyterlab-manager:saveWidgetState' }
+      { command: '@jupyter-widgets/jupyterlab-manager:saveNoWidgets' },
+      { command: '@jupyter-widgets/jupyterlab-manager:saveDisplayedWidgets' },
+      { command: '@jupyter-widgets/jupyterlab-manager:saveAllWidgets' }
     ]);
   }