From bd79a4d6076b6b1ff7b9c382d5de6e8093463f0d Mon Sep 17 00:00:00 2001
From: Edan Bainglass <edan.bainglass@gmail.com>
Date: Wed, 20 Nov 2024 09:47:30 +0000
Subject: [PATCH] Improve node view loading

---
 src/aiidalab_qe/app/result/__init__.py | 44 +++++++++++---------------
 1 file changed, 18 insertions(+), 26 deletions(-)

diff --git a/src/aiidalab_qe/app/result/__init__.py b/src/aiidalab_qe/app/result/__init__.py
index 1fd1463c2..ecf5be6ed 100644
--- a/src/aiidalab_qe/app/result/__init__.py
+++ b/src/aiidalab_qe/app/result/__init__.py
@@ -90,11 +90,7 @@ def render(self):
             "selected_nodes",
         )
 
-        self.container = ipw.VBox(
-            children=[
-                self.process_tree,
-            ],
-        )
+        self.node_view_container = ipw.VBox()
 
         self.process_monitor = ProcessMonitor(
             timeout=0.2,
@@ -115,7 +111,8 @@ def render(self):
                 ],
                 layout=ipw.Layout(margin="0 3px"),
             ),
-            self.container,
+            self.process_tree,
+            self.node_view_container,
         ]
 
         self.rendered = True
@@ -188,31 +185,26 @@ def _update_node_view(self, nodes, refresh=False):
         # check if the viewer is already added
         if node.uuid in self.node_views and not refresh:
             self.node_view = self.node_views[node.uuid]
-        else:
-            self.container.children = [
-                self.process_tree,
-                self.node_view_loading_message,
-            ]
-            if not isinstance(node, orm.WorkChainNode):
-                self.node_view = node_viewer(node)
-                self.node_views[node.uuid] = self.node_view
-            elif node.process_label == "QeAppWorkChain":
-                self._create_workchain_viewer(node)
-
-        self.container.children = [
-            self.process_tree,
-            self.node_view,
-        ]
-
-    def _create_workchain_viewer(self, node):
+        elif not isinstance(node, orm.WorkChainNode):
+            self.node_view_container.children = [self.node_view_loading_message]
+            self.node_view = node_viewer(node)
+            self.node_views[node.uuid] = self.node_view
+        elif node.process_label == "QeAppWorkChain":
+            self.node_view_container.children = [self.node_view_loading_message]
+            self.node_view = self._create_workchain_viewer(node)
+            self.node_views[node.uuid] = self.node_view
+
+        self.node_view_container.children = [self.node_view]
+
+    def _create_workchain_viewer(self, node: orm.WorkChainNode):
         model = WorkChainViewerModel()
         ipw.dlink(
             (self._model, "monitor_counter"),
             (model, "monitor_counter"),
         )
-        self.node_view: WorkChainViewer = node_viewer(node, model=model)
-        self.node_view.render()
-        self.node_views[node.uuid] = self.node_view
+        node_view: WorkChainViewer = node_viewer(node, model=model)  # type: ignore
+        node_view.render()
+        return node_view
 
     def _update_kill_button_layout(self):
         if not self.rendered: