Fix nested async/wait conversion

ibdafna · Oct 5, 2022 · c4f468e · c4f468e
1 parent d9007f0
commit c4f468e
Showing 1 changed file with 53 additions and 64 deletions.
diff --git a/packages/base/src/widget.ts b/packages/base/src/widget.ts
@@ -1,8 +1,8 @@
 // Copyright (c) Jupyter Development Team.
 // Distributed under the terms of the Modified BSD License.
 
-import * as utils from './utils';
 import * as backbonePatch from './backbone-patch';
+import * as utils from './utils';
 
 import * as Backbone from 'backbone';
 import $ from 'jquery';
@@ -13,15 +13,15 @@ import { JSONObject } from '@lumino/coreutils';
 
 import { Message, MessageLoop } from '@lumino/messaging';
 
-import { Widget, Panel } from '@lumino/widgets';
+import { Panel, Widget } from '@lumino/widgets';
 
 import { LayoutModel } from './widget_layout';
 
 import { StyleModel } from './widget_style';
 
 import { IWidgetManager } from './manager';
 
-import { IClassicComm, ICallbacks } from './services-shim';
+import { ICallbacks, IClassicComm } from './services-shim';
 
 import { JUPYTER_WIDGETS_VERSION } from './version';
 
@@ -218,25 +218,62 @@ export class WidgetModel extends Backbone.Model {
   /**
    * Handle incoming comm msg.
    */
-  _handle_comm_msg(msg: KernelMessage.ICommMsgMsg): Promise<void> {
+  async _handle_comm_msg(msg: KernelMessage.ICommMsgMsg): Promise<void> {
     const data = msg.content.data as any;
     const method = data.method;
     switch (method) {
       case 'update':
       case 'echo_update':
-        this.state_change = this.state_change
-          .then(() => {
-            return this._create_state_promise(msg);
-          })
-          .then((state) => {
-            this.set_state(state);
-          })
-          .catch(
-            utils.reject(
-              `Could not process update msg for model id: ${this.model_id}`,
-              true
-            )
+        try {
+          // this.state_change = this.state_change
+          this.state_change = await this.state_change;
+          const state: Dict<BufferJSON> = data.state;
+          const buffer_paths = data.buffer_paths ?? [];
+          const buffers = msg.buffers?.slice(0, buffer_paths.length) ?? [];
+          utils.put_buffers(state, buffer_paths, buffers);
+
+          if (msg.parent_header && method === 'echo_update') {
+            const msgId = (msg.parent_header as any).msg_id;
+            // we may have echos coming from other clients, we only care about
+            // dropping echos for which we expected a reply
+            const expectedEcho = Object.keys(state).filter((attrName) =>
+              this._expectedEchoMsgIds.has(attrName)
+            );
+            expectedEcho.forEach((attrName: string) => {
+              // Skip echo messages until we get the reply we are expecting.
+              const isOldMessage =
+                this._expectedEchoMsgIds.get(attrName) !== msgId;
+              if (isOldMessage) {
+                // Ignore an echo update that comes before our echo.
+                delete state[attrName];
+              } else {
+                // we got our echo confirmation, so stop looking for it
+                this._expectedEchoMsgIds.delete(attrName);
+                // Start accepting echo updates unless we plan to send out a new state soon
+                if (
+                  this._msg_buffer !== null &&
+                  Object.prototype.hasOwnProperty.call(
+                    this._msg_buffer,
+                    attrName
+                  )
+                ) {
+                  delete state[attrName];
+                }
+              }
+            });
+          }
+          const return_val = (this.constructor as typeof WidgetModel)._deserialize_state(
+            // Combine the state updates, with preference for kernel updates
+            state,
+            this.widget_manager
           );
+          this.set_state(await return_val);
+        } catch (error) {
+          utils.reject(
+            `Could not process update msg for model id: ${this.model_id}`,
+            true
+          )
+        }
         return this.state_change;
       case 'custom':
         this.trigger('msg:custom', data.content, msg.buffers);
@@ -245,54 +282,6 @@ export class WidgetModel extends Backbone.Model {
     return Promise.resolve();
   }
 
-  /**
-   * Handles echo update and returns deserialized state as a promise
-   * which we use to update this.state_change
-   */
-  async _create_state_promise(
-    msg: KernelMessage.ICommMsgMsg
-  ): Promise<utils.Dict<utils.BufferJSON>> {
-    const data = msg.content.data as any;
-    const method = data.method;
-
-    const state: Dict<BufferJSON> = data.state;
-    const buffer_paths = data.buffer_paths ?? [];
-    const buffers = msg.buffers?.slice(0, buffer_paths.length) ?? [];
-    utils.put_buffers(state, buffer_paths, buffers);
-
-    if (msg.parent_header && method === 'echo_update') {
-      const msgId = (msg.parent_header as any).msg_id;
-      // we may have echos coming from other clients, we only care about
-      // dropping echos for which we expected a reply
-      const expectedEcho = Object.keys(state).filter((attrName) =>
-        this._expectedEchoMsgIds.has(attrName)
-      );
-      expectedEcho.forEach((attrName: string) => {
-        // Skip echo messages until we get the reply we are expecting.
-        const isOldMessage = this._expectedEchoMsgIds.get(attrName) !== msgId;
-        if (isOldMessage) {
-          // Ignore an echo update that comes before our echo.
-          delete state[attrName];
-        } else {
-          // we got our echo confirmation, so stop looking for it
-          this._expectedEchoMsgIds.delete(attrName);
-          // Start accepting echo updates unless we plan to send out a new state soon
-          if (
-            this._msg_buffer !== null &&
-            Object.prototype.hasOwnProperty.call(this._msg_buffer, attrName)
-          ) {
-            delete state[attrName];
-          }
-        }
-      });
-    }
-    return await (this.constructor as typeof WidgetModel)._deserialize_state(
-      // Combine the state updates, with preference for kernel updates
-      state,
-      this.widget_manager
-    );
-  }
-
   /**
    * Handle when a widget is updated from the backend.
    *