Stop ES modules from waiting forever on modules that never load

ipyreact/module.py

@@ -1,6 +1,6 @@
 import warnings
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import traitlets
 from ipywidgets import DOMWidget
@@ -9,6 +9,11 @@
 from ._frontend import module_name, module_version
 
 _standard_dependencies: List[str] = []
+# the latest widget per name, a closed one provides nothing to depend on
+_module_widgets: Dict[str, "Module"] = {}
+# like solara, a name keeps the dependencies it was first defined with, so a
+# redefinition cannot pick up later modules and form a cycle (a->[b], b->[a])
+_module_dependencies: Dict[str, List[str]] = {}
 
 
 class Module(DOMWidget):
@@ -29,14 +34,26 @@ class Module(DOMWidget):
 
     @default("dependencies")
     def _default_dependencies(self):
-        return [k for k in get_module_names() if k != self.name]
+        return [k for k in get_module_names() if k != self.name and _is_live(k)]
+
+
+def _is_live(name):
+    widget = _module_widgets.get(name)
+    return widget is None or widget.comm is not None
 
 
 def get_module_names():
     return _standard_dependencies
 
 
-def define_module(name, module: Union[str, Path, None] = None, *, code: Optional[str] = None, url: Optional[str] = None):
+def define_module(
+    name,
+    module: Union[str, Path, None] = None,
+    *,
+    code: Optional[str] = None,
+    url: Optional[str] = None,
+    dependencies: Optional[List[str]] = None,
+):
     """Register a ES module under a name.
 
     Parameters
@@ -50,6 +67,10 @@ def define_module(name, module: Union[str, Path, None] = None, *, code: Optional
     url: str
         A url the module is served from (e.g. a bundle in the app's
         static dir)
+    dependencies: list of str
+        Module names to wait for before loading. Defaults to the
+        dependencies of the first definition of this name, else to the live
+        modules defined earlier in this process.
     """
     if sum(x is not None for x in (module, code, url)) != 1:
         raise TypeError("pass exactly one of module (a Path), code or url")
@@ -62,10 +83,26 @@ def define_module(name, module: Union[str, Path, None] = None, *, code: Optional
         )
         code = module
         module = None
-    _standard_dependencies.append(name)
-    if url is not None:
-        return Module(url=url, name=name)
-    if code is None:
+    if code is None and url is None:
         assert module is not None
         code = module.read_text(encoding="utf8")
-    return Module(code=code, name=name)
+    widget = _module_widgets.get(name)
+    if widget is not None and widget.comm is not None:
+        # redefining updates the live widget, like solara does: a second
+        # widget would get the modules defined after the first as
+        # dependencies, and wait for them forever (a->[b], b->[a])
+        with widget.hold_sync():
+            widget.url = url
+            widget.code = code or ""
+            if dependencies is not None:
+                widget.dependencies = dependencies
+        return widget
+    if dependencies is None:
+        dependencies = _module_dependencies.get(name)
+    if name not in _standard_dependencies:
+        _standard_dependencies.append(name)
+    kwargs = {} if dependencies is None else {"dependencies": dependencies}
+    widget = Module(code=code or "", url=url, name=name, **kwargs)
+    _module_dependencies.setdefault(name, widget.dependencies)
+    _module_widgets[name] = widget
+    return widget

ipyreact/widget.py

@@ -95,7 +95,7 @@ def _events(self, change):
     def _default_dependencies(self):
         import ipyreact.module
 
-        return ipyreact.module.get_module_names()
+        return [name for name in ipyreact.module.get_module_names() if ipyreact.module._is_live(name)]
 
 
 class ValueWidget(Widget, ValueWidgetClassic):

src/widget.tsx

@@ -75,11 +75,16 @@ function provideModule(moduleName: string, module: any) {
 function invalidateModule(moduleName: string) {
   // Next requestModule waits for a fresh provideModule. This is called
   // synchronously for both replacement models and trait updates so consumers
-  // created during hot reload never see the stale module.
-  delete modules[moduleName];
-  delete moduleFunctions[moduleName];
+  // created during hot reload never see the stale module. A pending entry is
+  // kept: consumers already waiting for it get the module the new load provides.
+  if (!moduleFunctions[moduleName]) {
+    delete modules[moduleName];
+  }
 }
 
+// per module name, so a slow load of old code cannot overwrite newer code
+const moduleGenerations: { [name: string]: number } = {};
+
 function requestModule(moduleName: string) {
   if (!modules[moduleName]) {
     modules[moduleName] = new Promise((resolve, reject) => {
@@ -346,6 +351,9 @@ export class Module extends WidgetModel {
   async addModule() {
     const code = this.get("code");
     let name = this.get("name");
+    const generation = (moduleGenerations[name] || 0) + 1;
+    moduleGenerations[name] = generation;
+    const isStale = () => moduleGenerations[name] !== generation;
     try {
       if (this.codeUrl && this.codeUrl.startsWith("blob:")) {
         URL.revokeObjectURL(this.codeUrl);
@@ -362,8 +370,14 @@ export class Module extends WidgetModel {
       await Promise.all(dependencies.map((x: any) => requestModule(x)));
       await ensureImportShimLoaded();
       await this.updateImportMap();
+      if (isStale()) {
+        return;
+      }
       this.set("status", "Loading module...");
       let module = await importShim(this.codeUrl!);
+      if (isStale()) {
+        return;
+      }
       try {
         // remapping an already-resolved specifier throws (hot reload in the
         // same page); the module registry is the source of truth, so only
@@ -375,6 +389,9 @@ export class Module extends WidgetModel {
       this.set("status", "Loaded module!");
       provideModule(name, module);
     } catch (e) {
+      if (isStale()) {
+        return;
+      }
       console.error(e);
       provideModule(name, e);
       this.set("status", "Error loading module: " + e);
@@ -653,7 +670,12 @@ export class ReactModel extends DOMWidgetModel {
           throw new Error(`Error loading module`);
         }
       } else {
-        module = await importShim(moduleName);
+        // a Module widget's name comes from the registry: the import map
+        // keeps the first url when the page's shim rejects map overrides
+        // (ipyvue >= 3 sets esmsInitOptions without mapOverrides)
+        module = moduleGenerations[moduleName]
+          ? await requestModule(moduleName)
+          : await importShim(moduleName);
         if (!module) {
           throw new Error(`no module found with name ${moduleName}`);
         }

tests/ui/module_hot_reload_test.py

@@ -45,3 +45,16 @@ def test_module_recreated_widget_hot_reload(solara_test, page_session: playwrigh
     ipyreact.define_module("recreate-module", code_template % 2)
     display(ipyreact.ValueWidget(_module="recreate-module", _type="Label"))
     page_session.locator(".hot-widget >> text=version 2").wait_for()
+
+
+def test_module_code_change_while_pending(solara_test, page_session: playwright.sync_api.Page):
+    # A consumer that waits for a module whose code changes before the first
+    # load finishes must get the new code, and not wait forever.
+    module = ipyreact.module.Module(name="pending-module", code=code_template % 1, dependencies=["pending-module-dep"])
+    display(ipyreact.ValueWidget(_module="pending-module", _type="Label", _dependencies=["pending-module"]))
+    display(module)
+    page_session.locator("text=pending-module status: Waiting for dependencies: pending-module-dep").wait_for()
+
+    module.code = code_template % 2
+    ipyreact.module.Module(name="pending-module-dep", code="export default 1", dependencies=[])
+    page_session.locator(".hot-widget >> text=version 2").wait_for()

tests/unit/define_module_test.py

@@ -35,3 +35,47 @@ def test_define_module_exactly_one():
         ipyreact.define_module("t-none")
     with pytest.raises(TypeError, match="exactly one"):
         ipyreact.define_module("t-both", Path("x"), code="y")
+
+
+def test_define_module_dependencies():
+    ipyreact.define_module("t-deps-a", code="export default 1")
+    module = ipyreact.define_module("t-deps-b", code="export default 2", dependencies=["t-deps-a"])
+    assert module.dependencies == ["t-deps-a"]
+    module = ipyreact.define_module("t-deps-c", code="export default 3", dependencies=[])
+    assert module.dependencies == []
+
+
+def test_define_module_skips_closed_dependencies():
+    ipyreact.define_module("t-closed-a", code="export default 1").close()
+    module = ipyreact.define_module("t-closed-b", code="export default 2")
+    assert "t-closed-a" not in module.dependencies
+
+
+def test_define_module_redefine_keeps_dependencies():
+    a = ipyreact.define_module("t-redef-a", code="export default 1")
+    b = ipyreact.define_module("t-redef-b", code="export default 2")
+    assert "t-redef-a" in b.dependencies
+    # re-running the cell updates the live widgets, without a cycle
+    assert ipyreact.define_module("t-redef-a", code="export default 3") is a
+    assert a.code == "export default 3"
+    assert "t-redef-b" not in a.dependencies
+    assert ipyreact.define_module("t-redef-b", url="/static/b.mjs") is b
+    assert b.url == "/static/b.mjs"
+    assert b.code == ""
+    # a closed widget is recreated with the dependencies of its first definition
+    a.close()
+    a2 = ipyreact.define_module("t-redef-a", code="export default 4")
+    assert a2 is not a
+    assert a2.dependencies == a.dependencies
+    assert ipyreact.define_module("t-redef-a", code="export default 5", dependencies=[]).dependencies == []
+
+
+def test_widget_default_dependencies_skip_closed_modules():
+    closed = ipyreact.define_module("t-consumer-closed", code="export default 1")
+    ipyreact.define_module("t-consumer-live", code="export default 2")
+    closed.close()
+
+    deps = ipyreact.Widget()._dependencies
+
+    assert "t-consumer-live" in deps
+    assert "t-consumer-closed" not in deps