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

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
+): void {
+  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 IWidgetRelations {
+  widget: WidgetModel;
+  children: Set<WidgetModel>;
+  parents: Set<WidgetModel>;
+}
+
+function _removeConnected(
+  widget_id: string,
+  relations: Map<string, IWidgetRelations>
+): void {
+  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>
+): Set<WidgetModel> {
+  const relations = new Map<string, IWidgetRelations>();
+  // builds up a 'graph'
+  allWidgets.forEach((widget, widget_id) => {
+    const relation: IWidgetRelations = {
+      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;
   }