feat: always save state of widgets, and only those required

The default is to save widgets state in the notebook, but to avoid
it being too large, we only save those that are needed to display
the visible ones (and thus find all widgets that are connected).

jupyterlab_widgets/src/manager.ts

@@ -36,6 +36,9 @@ import { ISignal, Signal } from '@lumino/signaling';
 import { valid } from 'semver';
 
 import { SemVerCache } from './semvercache';
+import { toArray } from '@lumino/algorithm';
+import { ICodeCellModel } from '@jupyterlab/cells';
+import { findConnectedWidgets } from '@jupyter-widgets/base';
 
 /**
  * The mime type for a widget view.
@@ -340,30 +343,6 @@ export abstract class LabWidgetManager extends ManagerBase
     return modelPromise;
   }
 
-  /**
-   * Register a widget model.
-   */
-  register_model(model_id: string, modelPromise: Promise<WidgetModel>): void {
-    super.register_model(model_id, modelPromise);
-
-    // Update the synchronous model map
-    modelPromise.then(model => {
-      this._modelsSync.set(model_id, model);
-      model.once('comm:close', () => {
-        this._modelsSync.delete(model_id);
-      });
-    });
-  }
-
-  /**
-   * Close all widgets and empty the widget state.
-   * @return Promise that resolves when the widget state is cleared.
-   */
-  async clear_state(): Promise<void> {
-    await super.clear_state();
-    this._modelsSync = new Map();
-  }
-
   /**
    * Synchronously get the state of the live widgets in the widget manager.
    *
@@ -374,10 +353,17 @@ export abstract class LabWidgetManager extends ManagerBase
    * @returns A state dictionary
    */
   get_state_sync(options: IStateOptions = {}): ReadonlyPartialJSONValue {
-    const models = [];
-    for (const model of this._modelsSync.values()) {
-      if (model.comm_live) {
-        models.push(model);
+    let models: Array<WidgetModel> = [];
+    if (options.visibleWidgets) {
+      models = [
+        ...findConnectedWidgets(options.visibleWidgets, this.get_models_sync())
+      ];
+      console.log('Saving only visible widgets: ', models);
+    } else {
+      for (const model of this.get_models_sync().values()) {
+        if (model.comm_live) {
+          models.push(model);
+        }
       }
     }
     return serialize_state(models, options);
@@ -403,7 +389,6 @@ export abstract class LabWidgetManager extends ManagerBase
 
   private _commRegistration: IDisposable;
 
-  private _modelsSync = new Map<string, WidgetModel>();
   private _onUnhandledIOPubMessage = new Signal<
     this,
     KernelMessage.IIOPubMessage
@@ -532,7 +517,24 @@ export class WidgetManager extends LabWidgetManager {
    * Save the widget state to the context model.
    */
   private _saveState(): void {
-    const state = this.get_state_sync({ drop_defaults: true });
+    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 state = this.get_state_sync({
+      drop_defaults: true,
+      visibleWidgets: visibleWidgets
+    });
     this._context.model.metadata.set('widgets', {
       'application/vnd.jupyter.widget-state+json': state
     });

jupyterlab_widgets/src/plugin.ts

@@ -53,7 +53,7 @@ const WIDGET_REGISTRY: base.IWidgetRegistryData[] = [];
 /**
  * The cached settings.
  */
-const SETTINGS: WidgetManager.Settings = { saveState: false };
+const SETTINGS: WidgetManager.Settings = { saveState: true };
 
 /**
  * Iterate through all widget renderers in a notebook.

packages/base-manager/src/manager-base.ts

@@ -142,6 +142,21 @@ export abstract class ManagerBase implements IWidgetManager {
     return this._models[model_id];
   }
 
+  /**
+   * Return a Map with all widgets that are created, note that this misses widgets
+   * that are in the process of being created (e.g. we may be in the process of fetching
+   * libraries)
+   */
+
+  get_models_sync(): Map<string, WidgetModel> {
+    // since map is mutable, return a copy
+    const copy = new Map<string, WidgetModel>();
+    this._modelsSync.forEach((value, key) => {
+      copy.set(key, value);
+    });
+    return copy;
+  }
+
   /**
    * Handle when a comm is opened.
    */
@@ -247,6 +262,7 @@ export abstract class ManagerBase implements IWidgetManager {
   register_model(model_id: string, modelPromise: Promise<WidgetModel>): void {
     this._models[model_id] = modelPromise;
     modelPromise.then(model => {
+      this._modelsSync.set(model_id, model);
       model.once('comm:close', () => {
         delete this._models[model_id];
       });
@@ -337,6 +353,7 @@ export abstract class ManagerBase implements IWidgetManager {
     return resolvePromisesDict(this._models).then(models => {
       Object.keys(models).forEach(id => models[id].close());
       this._models = Object.create(null);
+      this._modelsSync = new Map();
     });
   }
 
@@ -524,6 +541,7 @@ export abstract class ManagerBase implements IWidgetManager {
   private _models: { [key: string]: Promise<WidgetModel> } = Object.create(
     null
   );
+  private _modelsSync = new Map<string, WidgetModel>();
 }
 
 export interface IStateOptions {
@@ -533,6 +551,7 @@ export interface IStateOptions {
    * @default false
    */
   drop_defaults?: boolean;
+  visibleWidgets?: Array<string>;
 }
 
 /**

packages/base/src/utils.ts

@@ -4,6 +4,7 @@
 import { JSONObject, JSONValue, UUID, JSONExt } from '@lumino/coreutils';
 
 import _isEqual from 'lodash/isEqual';
+import { WidgetModel } from './widget';
 
 /**
  * Find all strings in the first argument that are not in the second.
@@ -230,3 +231,102 @@ export function remove_buffers(
   const new_state = remove(state, []) as JSONObject;
   return { state: new_state, buffers: buffers, buffer_paths: buffer_paths };
 }
+
+export function findWidgetChildren(
+  widget: WidgetModel,
+  recursive = false
+): Set<WidgetModel> {
+  const children = new Set<WidgetModel>();
+  _findWidgetChildren(widget, children, recursive ? Infinity : 1);
+  children.delete(widget);
+  return children;
+}
+
+function _findWidgetChildren(
+  obj: any,
+  found: Set<WidgetModel>,
+  levels: number
+) {
+  if (obj instanceof WidgetModel) {
+    const widget: WidgetModel = obj;
+    // stop collecting if already found
+    if (found.has(widget)) {
+      return;
+    }
+    found.add(widget);
+    if (levels >= 1) {
+      for (const name in widget.attributes) {
+        const value = widget.attributes[name];
+        _findWidgetChildren(value, found, levels - 1);
+      }
+    }
+  } else if (Array.isArray(obj)) {
+    for (const value of obj) {
+      _findWidgetChildren(value, found, levels);
+    }
+  } else if (isObject(obj)) {
+    for (const key in obj) {
+      if (Object.prototype.hasOwnProperty.call(obj, key)) {
+        const value = obj[key];
+        _findWidgetChildren(value, found, levels);
+      }
+    }
+  }
+}
+
+interface WidgetRelations {
+  widget: WidgetModel;
+  children: Set<WidgetModel>;
+  parents: Set<WidgetModel>;
+}
+
+function _removeConnected(
+  widget_id: String,
+  relations: Map<String, WidgetRelations>
+) {
+  const relation = relations.get(widget_id);
+  if (relation) {
+    relations.delete(widget_id);
+    relation.children.forEach(child => {
+      _removeConnected(child.model_id, relations);
+    });
+    relation.parents.forEach(parent => {
+      _removeConnected(parent.model_id, relations);
+    });
+  }
+}
+
+export function findConnectedWidgets(
+  widgets: Array<String>,
+  allWidgets: Map<String, WidgetModel>
+) {
+  const relations = new Map<String, WidgetRelations>();
+  // builds up a 'graph'
+  allWidgets.forEach((widget, widget_id) => {
+    const relation: WidgetRelations = {
+      widget,
+      children: widget.getChildren(),
+      parents: new Set()
+    };
+    relations.set(widget_id, relation);
+  });
+  // find all parents
+  allWidgets.forEach((widget, widget_id) => {
+    relations.get(widget_id)?.children.forEach(child => {
+      relations.get(child.model_id)?.parents.add(widget);
+    });
+  });
+  // It is easier to think of solving the inverse problem. So we delete all widgets that
+  // we connect to from relations
+  widgets.forEach(model_id => {
+    _removeConnected(model_id, relations);
+  });
+  // Then the connected widgets, are all widgets that are not anymore in the relations map
+  const connected = new Set<WidgetModel>();
+  allWidgets.forEach((widget, widget_id) => {
+    if (!relations.has(widget_id)) {
+      connected.add(widget);
+    }
+  });
+  return connected;
+}

packages/base/src/widget.ts

@@ -25,7 +25,7 @@ import { IClassicComm, ICallbacks } from './services-shim';
 
 import { JUPYTER_WIDGETS_VERSION } from './version';
 
-import { Dict } from './utils';
+import { Dict, findWidgetChildren } from './utils';
 
 import { KernelMessage } from '@jupyterlab/services';
 
@@ -143,6 +143,15 @@ export class WidgetModel extends Backbone.Model {
     }
   }
 
+  /**
+   * Get a set of all children widget. The default implementation will look for all
+   * attributes, and resursively inspect Arrays and Objects. If this is not sufficient,
+   * library authors can override this method.
+   */
+  getChildren(): Set<WidgetModel> {
+    return findWidgetChildren(this);
+  }
+
   get comm_live(): boolean {
     return this._comm_live;
   }