refactor: Refactor model initialization logic

src/index.tsx

@@ -1,6 +1,6 @@
 import { createRender, useModel, useModelState } from "@anywidget/react";
 import type { Initialize, Render } from "@anywidget/types";
-import { MapViewState, PickingInfo, type Layer } from "@deck.gl/core";
+import { MapViewState, PickingInfo } from "@deck.gl/core";
 import { DeckGLRef } from "@deck.gl/react";
 import type { IWidgetManager, WidgetModel } from "@jupyter-widgets/base";
 import { NextUIProvider } from "@nextui-org/react";
@@ -10,7 +10,11 @@ import { useCallback, useEffect, useRef, useState } from "react";
 import { v4 as uuidv4 } from "uuid";
 
 import { flyTo } from "./actions/fly-to.js";
-import { BaseLayerModel, initializeLayer } from "./model/index.js";
+import {
+  initializeLayer,
+  type BaseLayerModel,
+  initializeChildModels,
+} from "./model/index.js";
 import { initParquetWasm } from "./parquet.js";
 import DeckFirstRenderer from "./renderers/deck-first.js";
 import OverlayRenderer from "./renderers/overlay.js";
@@ -20,7 +24,7 @@ import { useViewStateDebounced } from "./state";
 import Toolbar from "./toolbar.js";
 import { getTooltip } from "./tooltip/index.js";
 import { Message } from "./types.js";
-import { isDefined, loadChildModels } from "./util.js";
+import { isDefined } from "./util.js";
 import { MachineContext, MachineProvider } from "./xstate";
 import * as selectors from "./xstate/selectors";
 
@@ -40,40 +44,6 @@ const DEFAULT_INITIAL_VIEW_STATE = {
 const DEFAULT_MAP_STYLE =
   "https://basemaps.cartocdn.com/gl/positron-nolabels-gl-style/style.json";
 
-async function getChildModelState(
-  childModels: WidgetModel[],
-  childLayerIds: string[],
-  previousSubModelState: Record<string, BaseLayerModel>,
-  setStateCounter: React.Dispatch<React.SetStateAction<Date>>,
-): Promise<Record<string, BaseLayerModel>> {
-  const newSubModelState: Record<string, BaseLayerModel> = {};
-  const updateStateCallback = () => setStateCounter(new Date());
-
-  for (let i = 0; i < childLayerIds.length; i++) {
-    const childLayerId = childLayerIds[i];
-    const childModel = childModels[i];
-
-    // If the layer existed previously, copy its model without constructing
-    // a new one
-    if (childLayerId in previousSubModelState) {
-      // pop from old state
-      newSubModelState[childLayerId] = previousSubModelState[childLayerId];
-      delete previousSubModelState[childLayerId];
-      continue;
-    }
-
-    const childLayer = await initializeLayer(childModel, updateStateCallback);
-    newSubModelState[childLayerId] = childLayer;
-  }
-
-  // finalize models that were deleted
-  for (const previousChildModel of Object.values(previousSubModelState)) {
-    previousChildModel.finalize();
-  }
-
-  return newSubModelState;
-}
-
 function App() {
   const actorRef = MachineContext.useActorRef();
   const isDrawingBBoxSelection = MachineContext.useSelector(
@@ -144,7 +114,7 @@ function App() {
   });
 
   const [mapId] = useState(uuidv4());
-  const [subModelState, setSubModelState] = useState<
+  const [layersState, setLayersState] = useState<
     Record<string, BaseLayerModel>
   >({});
 
@@ -153,22 +123,20 @@ function App() {
   // Fake state just to get react to re-render when a model callback is called
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   const [stateCounter, setStateCounter] = useState<Date>(new Date());
+  const updateStateCallback = () => setStateCounter(new Date());
 
   useEffect(() => {
     const loadAndUpdateLayers = async () => {
       try {
-        const childModels = await loadChildModels(
+        const layerModels = await initializeChildModels<BaseLayerModel>(
           model.widget_manager as IWidgetManager,
           childLayerIds,
+          layersState,
+          async (model: WidgetModel) =>
+            initializeLayer(model, updateStateCallback),
         );
 
-        const newSubModelState = await getChildModelState(
-          childModels,
-          childLayerIds,
-          subModelState,
-          setStateCounter,
-        );
-        setSubModelState(newSubModelState);
+        setLayersState(layerModels);
 
         if (!isDrawingBBoxSelection) {
           // Note: selected_bounds is a property of the **Map**. In the future,
@@ -189,15 +157,9 @@ function App() {
     loadAndUpdateLayers();
   }, [childLayerIds, bboxSelectBounds, isDrawingBBoxSelection]);
 
-  const layers: Layer[] = [];
-  for (const subModel of Object.values(subModelState)) {
-    const newLayers = subModel.render();
-    if (Array.isArray(newLayers)) {
-      layers.push(...newLayers);
-    } else {
-      layers.push(newLayers);
-    }
-  }
+  const layers = Object.values(layersState).flatMap((layerModel) =>
+    layerModel.render(),
+  );
 
   const onMapClickHandler = useCallback((info: PickingInfo) => {
     // We added this flag to prevent the hover event from firing after a

src/model/base-layer.ts

@@ -8,10 +8,11 @@ import type {
 } from "@deck.gl/core";
 import type { WidgetModel } from "@jupyter-widgets/base";
 
-import { isDefined, loadChildModels } from "../util.js";
 import { BaseModel } from "./base.js";
 import { initializeExtension } from "./extension.js";
 import type { BaseExtensionModel } from "./extension.js";
+import { initializeChildModels } from "./initialize.js";
+import { isDefined } from "../util.js";
 
 export abstract class BaseLayerModel extends BaseModel {
   protected pickable: LayerProps["pickable"];
@@ -20,7 +21,7 @@ export abstract class BaseLayerModel extends BaseModel {
   protected autoHighlight: LayerProps["autoHighlight"];
   protected highlightColor: LayerProps["highlightColor"];
 
-  protected extensions: BaseExtensionModel[];
+  protected extensions: Record<string, BaseExtensionModel>;
 
   /** Names of additional layer properties that are dynamically added by
    * extensions and should be rendered with layer attributes.
@@ -37,15 +38,15 @@ export abstract class BaseLayerModel extends BaseModel {
     this.initRegularAttribute("highlight_color", "highlightColor");
     this.initRegularAttribute("selected_bounds", "selectedBounds");
 
-    this.extensions = [];
+    this.extensions = {};
   }
 
   async loadSubModels() {
     await this.initLayerExtensions();
   }
 
   extensionInstances(): LayerExtension[] {
-    return this.extensions
+    return Object.values(this.extensions)
       .map((extension) => extension.extensionInstance())
       .filter((extensionInstance) => extensionInstance !== null);
   }
@@ -108,29 +109,19 @@ export abstract class BaseLayerModel extends BaseModel {
   // experimental
   async initLayerExtensions() {
     const initExtensionsCallback = async () => {
-      const childModelIds = this.model.get("extensions");
-      if (!childModelIds) {
-        this.extensions = [];
-        return;
-      }
+      const extensionModelIds = this.model.get("extensions");
 
-      const childModels = await loadChildModels(
+      const extensionModels = await initializeChildModels<BaseExtensionModel>(
         this.model.widget_manager,
-        childModelIds,
+        extensionModelIds,
+        this.extensions,
+        async (childModel: WidgetModel) =>
+          initializeExtension(childModel, this, this.updateStateCallback),
       );
 
-      const extensions: BaseExtensionModel[] = [];
-      for (const childModel of childModels) {
-        const extension = await initializeExtension(
-          childModel,
-          this,
-          this.updateStateCallback,
-        );
-        extensions.push(extension);
-      }
-
-      this.extensions = extensions;
+      this.extensions = extensionModels;
     };
+
     await initExtensionsCallback();
 
     // Remove all existing change callbacks for this attribute

src/model/index.ts

@@ -13,3 +13,4 @@ export {
   PathStyleExtension,
   initializeExtension,
 } from "./extension.js";
+export { initializeChildModels } from "./initialize.js";

src/model/initialize.ts

@@ -0,0 +1,75 @@
+import type { IWidgetManager, WidgetModel } from "@jupyter-widgets/base";
+
+import type { BaseModel } from "./base";
+
+// childModelId is of the form "IPY_MODEL_<identifier>"
+// https://github.com/jupyter-widgets/ipywidgets/blob/8.1.7/packages/schema/jupyterwidgetmodels.v8.md
+const IPY_MODEL_ = "IPY_MODEL_";
+
+// We need to slice off the "IPY_MODEL_" prefix to get the actual model ID
+const IPY_MODEL_LENGTH = IPY_MODEL_.length;
+
+/**
+ * Load and initialize the child models of this model.
+ *
+ * @param widget_manager The widget manager used to load the models.
+ * @param childModelIds The model IDs of the child models to load.
+ * @param previousSubModelState Any previously loaded child models. Models that are still present will be reused. Any reactivity on _those models_ will be handled separately.
+ * @param initializer A function that takes a WidgetModel and returns an initialized model of type T.
+ *
+ * @return A promise that resolves to a mapping from model ID to initialized model.
+ */
+export async function initializeChildModels<T extends BaseModel>(
+  widget_manager: IWidgetManager,
+  childModelIds: string[],
+  previousSubModelState: Record<string, T>,
+  initializer: (model: WidgetModel) => Promise<T>,
+): Promise<Record<string, T>> {
+  const childModels = await loadModels(widget_manager, childModelIds);
+
+  const newSubModelState: Record<string, T> = {};
+
+  for (const [childModelId, childModel] of Object.entries(childModels)) {
+    // If the layer existed previously, copy its model without constructing
+    // a new one
+    if (childModelId in previousSubModelState) {
+      // reuse existing model and remove from old state
+      newSubModelState[childModelId] = previousSubModelState[childModelId];
+      delete previousSubModelState[childModelId];
+    } else {
+      // TODO: should we be using Promise.all here?
+      newSubModelState[childModelId] = await initializer(childModel);
+    }
+  }
+
+  // finalize models that were deleted
+  for (const previousChildModel of Object.values(previousSubModelState)) {
+    previousChildModel.finalize();
+  }
+
+  return newSubModelState;
+}
+
+/**
+ * Load and resolve other widget models.
+ *
+ * Loading of models is asynchronous; we load all models in parallel.
+ */
+async function loadModels(
+  widget_manager: IWidgetManager,
+  childModelIds: string[],
+): Promise<Record<string, WidgetModel>> {
+  const promises = childModelIds.map((childModelId) => {
+    // We need to slice off the "IPY_MODEL_" prefix to get the actual model ID
+    const modelId = childModelId.slice(IPY_MODEL_LENGTH);
+    return widget_manager.get_model(modelId);
+  });
+
+  const models = await Promise.all(promises);
+
+  const returnedModels: Record<string, WidgetModel> = {};
+  for (let i = 0; i < childModelIds.length; i++) {
+    returnedModels[childModelIds[i]] = models[i];
+  }
+  return returnedModels;
+}

src/util.ts

@@ -1,21 +1,3 @@
-import type { IWidgetManager, WidgetModel } from "@jupyter-widgets/base";
-
-/**
- * Load the child models of this model
- */
-export async function loadChildModels(
-  widget_manager: IWidgetManager,
-  childLayerIds: string[],
-): Promise<WidgetModel[]> {
-  const promises: Promise<WidgetModel>[] = [];
-  for (const childLayerId of childLayerIds) {
-    promises.push(
-      widget_manager.get_model(childLayerId.slice("IPY_MODEL_".length)),
-    );
-  }
-  return await Promise.all(promises);
-}
-
 /** Check for null and undefined */
 // https://stackoverflow.com/a/52097445
 export function isDefined<T>(value: T | undefined | null): value is T {