make effects run just before and after full render

Author: rmorshea

docs/source/examples/snake_game.py

@@ -108,7 +108,14 @@ async def animate():
 def use_async_effect(function):
     def ensure_effect_future():
         future = asyncio.ensure_future(function())
-        return future.cancel
+
+        def cleanup():
+            if future.done():
+                future.result()
+            else:
+                future.cancel()
+
+        return cleanup
 
     idom.hooks.use_effect(ensure_effect_future)
 
@@ -120,7 +127,7 @@ async def interval():
         await asyncio.sleep(rate - (time.time() - usage_time.current))
         usage_time.current = time.time()
 
-    return interval()
+    return asyncio.ensure_future(interval())
 
 
 def use_snake_food(grid_size, current_snake):

idom/core/hooks.py

@@ -14,6 +14,8 @@
     overload,
 )
 
+from loguru import logger
+
 from idom.utils import Ref
 
 from .element import AbstractElement
@@ -373,15 +375,24 @@ def element_will_render(self) -> None:
         self._is_rendering = True
 
         for effect in self._event_effects.will_render:
-            effect()
+            try:
+                effect()
+            except Exception:
+                msg = f"Pre-render effect {effect} failed for {self.element}"
+                logger.exception(msg)
 
         self._event_effects.will_render.clear()
         self._event_effects.will_unmount.clear()
 
     def element_did_render(self) -> None:
         """The element completed a render"""
         for effect in self._event_effects.did_render:
-            effect()
+            try:
+                effect()
+            except Exception:
+                msg = f"Post-render effect {effect} failed for {self.element}"
+                logger.exception(msg)
+
         self._event_effects.did_render.clear()
 
         self._is_rendering = False
@@ -393,7 +404,12 @@ def element_did_render(self) -> None:
     def element_will_unmount(self) -> None:
         """The element is about to be removed from the layout"""
         for effect in self._event_effects.will_unmount:
-            effect()
+            try:
+                effect()
+            except Exception:
+                msg = f"Pre-unmount effect {effect} failed for {self.element}"
+                logger.exception(msg)
+
         self._event_effects.will_unmount.clear()
 
     def set_current(self) -> None:

idom/core/layout.py

@@ -101,22 +101,36 @@ async def _element_states(self) -> AsyncIterator[ElementState]:
         try:
             yield {self._root.id: root_element_state}
         finally:
-            self._unmount_element_state(root_element_state)
+            self._delete_element_state(root_element_state)
 
     async def _create_layout_update(self, element: AbstractElement) -> LayoutUpdate:
         element_state = self._element_states[element.id]
+
+        element_state.life_cycle_hook.element_will_render()
+
+        for state in self._iter_element_states_from_root(
+            element_state,
+            include_root=False,
+        ):
+            state.life_cycle_hook.element_will_unmount()
+
+        self._clear_element_state_event_handlers(element_state)
+        self._delete_element_state_children(element_state)
+
         old_model = element_state.model.copy()  # we copy because it will be mutated
         new_model = await self._render_element(element_state)
         changes = make_patch(old_model, new_model).patch
+
+        for state in self._iter_element_states_from_root(
+            element_state,
+            include_root=True,
+        ):
+            state.life_cycle_hook.element_did_render()
+
         return LayoutUpdate(path=element_state.path, changes=changes)
 
     async def _render_element(self, element_state: ElementState) -> Dict[str, Any]:
         try:
-            element_state.life_cycle_hook.element_will_render()
-
-            self._clear_element_state_event_handlers(element_state)
-            self._unmount_element_state_children(element_state)
-
             # BUG: https://github.com/python/mypy/issues/9256
             raw_model = await _render_with_life_cycle_hook(element_state)  # type: ignore
 
@@ -126,8 +140,6 @@ async def _render_element(self, element_state: ElementState) -> Dict[str, Any]:
             resolved_model = await self._render_model(element_state, raw_model)
             element_state.model.clear()
             element_state.model.update(resolved_model)
-
-            element_state.life_cycle_hook.element_did_render()
         except Exception:
             logger.exception(f"Failed to render {element_state.element_obj}")
 
@@ -206,24 +218,43 @@ def _create_element_state(
 
     def _reset_element_state(self, element_state: ElementState) -> None:
         self._clear_element_state_event_handlers(element_state)
-        self._unmount_element_state_children(element_state)
+        self._delete_element_state_children(element_state)
 
-    def _unmount_element_state(self, element_state: ElementState) -> None:
-        element_state.life_cycle_hook.element_will_unmount()
+    def _delete_element_state(self, element_state: ElementState) -> None:
         self._clear_element_state_event_handlers(element_state)
-        self._unmount_element_state_children(element_state)
+        self._delete_element_state_children(element_state)
         del self._element_states[element_state.element_obj.id]
 
     def _clear_element_state_event_handlers(self, element_state: ElementState) -> None:
         for handler_id in element_state.event_handler_ids:
             del self._event_handlers[handler_id]
         element_state.event_handler_ids.clear()
 
-    def _unmount_element_state_children(self, element_state: ElementState) -> None:
+    def _delete_element_state_children(self, element_state: ElementState) -> None:
         for e_id in element_state.child_elements_ids:
-            self._unmount_element_state(self._element_states[e_id])
+            self._delete_element_state(self._element_states[e_id])
         element_state.child_elements_ids.clear()
 
+    def _iter_element_states_from_root(
+        self,
+        root_element_state: ElementState,
+        include_root: bool,
+    ) -> Iterator[ElementState]:
+        if include_root:
+            pending = [root_element_state]
+        else:
+            pending = [
+                self._element_states[i] for i in root_element_state.child_elements_ids
+            ]
+
+        while pending:
+            visited_element_state = pending.pop(0)
+            yield visited_element_state
+            pending.extend(
+                self._element_states[i]
+                for i in visited_element_state.child_elements_ids
+            )
+
 
 @coroutine
 def _render_with_life_cycle_hook(element_state: ElementState) -> Iterator[None]: