fix remove freeze

Author: willmcgugan

src/textual/app.py

@@ -39,6 +39,7 @@
 
 from . import Logger, LogGroup, LogVerbosity, actions, events, log, messages
 from ._animator import DEFAULT_EASING, Animatable, Animator, EasingFunction
+from .await_remove import AwaitRemove
 from ._ansi_sequences import SYNC_END, SYNC_START
 from ._callback import invoke
 from ._context import active_app
@@ -353,6 +354,7 @@ def __init__(
             else None
         )
         self._screenshot: str | None = None
+        self._dom_lock = asyncio.Lock()
 
     @property
     def return_value(self) -> ReturnType | None:
@@ -1936,6 +1938,42 @@ def _walk_children(self, root: Widget) -> Iterable[list[Widget]]:
             for child in widget.children:
                 push(child)
 
+    def _remove_nodes(self, widgets: list[Widget]) -> AwaitRemove:
+        """Remove nodes from DOM, and return an awaitable that awaits cleanup.
+
+        Args:
+            widgets (list[Widget]): List of nodes to remvoe.
+
+        Returns:
+            AwaitRemove: Awaitable that returns when the nodes have been fully removed.
+        """
+
+        async def remove_task(
+            widgets: list[Widget], finished_event: asyncio.Event
+        ) -> None:
+            try:
+                await self._prune_nodes(widgets)
+            finally:
+                finished_event.set()
+
+        removed_widgets = self._detach_from_dom(widgets)
+        self.refresh(layout=True)
+
+        finished_event = asyncio.Event()
+        asyncio.create_task(remove_task(removed_widgets, finished_event))
+
+        return AwaitRemove(finished_event)
+
+    async def _prune_nodes(self, widgets: list[Widget]) -> None:
+        """Remove nodes and children.
+
+        Args:
+            widgets (Widget): _description_
+        """
+        async with self._dom_lock:
+            for widget in widgets:
+                await self._prune_node(widget)
+
     async def _prune_node(self, root: Widget) -> None:
         """Remove a node and its children. Children are removed before parents.
 

src/textual/css/query.py

@@ -356,16 +356,9 @@ def remove(self) -> AwaitRemove:
         Returns:
             AwaitRemove: An awaitable object that waits for the widgets to be removed.
         """
-        prune_finished_event = asyncio.Event()
         app = active_app.get()
-        app.post_message_no_wait(
-            events.Prune(
-                app,
-                widgets=app._detach_from_dom(list(self)),
-                finished_flag=prune_finished_event,
-            )
-        )
-        return AwaitRemove(prune_finished_event)
+        await_remove = app._remove_nodes(list(self))
+        return await_remove
 
     def set_styles(
         self, css: str | None = None, **update_styles

src/textual/events.py

@@ -127,28 +127,6 @@ class Unmount(Mount, bubble=False, verbose=False):
     """Sent when a widget is unmounted and may not longer receive messages."""
 
 
-class Prune(Event, bubble=False):
-    """Sent to the app to ask it to prune one or more widgets from the DOM.
-
-    Attributes:
-        widgets (list[Widgets]): The list of widgets to prune.
-        finished_flag (asyncio.Event): An asyncio Event to that will be flagged when the prune is done.
-    """
-
-    def __init__(
-        self, sender: MessageTarget, widgets: list[Widget], finished_flag: asyncio.Event
-    ) -> None:
-        """Initialise the event.
-
-        Args:
-            widgets (list[Widgets]): The list of widgets to prune.
-            finished_flag (asyncio.Event): An asyncio Event to that will be flagged when the prune is done.
-        """
-        super().__init__(sender)
-        self.finished_flag = finished_flag
-        self.widgets = widgets
-
-
 class Show(Event, bubble=False):
     """Sent when a widget has become visible."""
 

src/textual/screen.py

@@ -310,18 +310,19 @@ async def _on_idle(self, event: events.Idle) -> None:
         # Check for any widgets marked as 'dirty' (needs a repaint)
         event.prevent_default()
 
-        if self.is_current:
-            if self._layout_required:
-                self._refresh_layout()
-                self._layout_required = False
-                self._dirty_widgets.clear()
-            if self._repaint_required:
-                self._dirty_widgets.clear()
-                self._dirty_widgets.add(self)
-                self._repaint_required = False
-
-            if self._dirty_widgets:
-                self.update_timer.resume()
+        async with self.app._dom_lock:
+            if self.is_current:
+                if self._layout_required:
+                    self._refresh_layout()
+                    self._layout_required = False
+                    self._dirty_widgets.clear()
+                if self._repaint_required:
+                    self._dirty_widgets.clear()
+                    self._dirty_widgets.add(self)
+                    self._repaint_required = False
+
+                if self._dirty_widgets:
+                    self.update_timer.resume()
 
         # The Screen is idle - a good opportunity to invoke the scheduled callbacks
         await self._invoke_and_clear_callbacks()

src/textual/widget.py

@@ -523,15 +523,19 @@ def mount(
 
         # Decide the final resting place depending on what we've been asked
         # to do.
+        insert_before: int | None = None
+        insert_after: int | None = None
         if before is not None:
-            parent, before = self._find_mount_point(before)
+            parent, insert_before = self._find_mount_point(before)
         elif after is not None:
-            parent, after = self._find_mount_point(after)
+            parent, insert_after = self._find_mount_point(after)
         else:
             parent = self
 
         return AwaitMount(
-            self.app._register(parent, *widgets, before=before, after=after)
+            self.app._register(
+                parent, *widgets, before=insert_before, after=insert_after
+            )
         )
 
     def move_child(
@@ -697,7 +701,6 @@ def get_content_height(self, container: Size, viewport: Size, width: int) -> int
         Returns:
             int: The height of the content.
         """
-
         if self.is_container:
             assert self._layout is not None
             height = (
@@ -2114,15 +2117,9 @@ def remove(self) -> AwaitRemove:
         Returns:
             AwaitRemove: An awaitable object that waits for the widget to be removed.
         """
-        prune_finished_event = AsyncEvent()
-        self.app.post_message_no_wait(
-            events.Prune(
-                self,
-                widgets=self.app._detach_from_dom([self]),
-                finished_flag=prune_finished_event,
-            )
-        )
-        return AwaitRemove(prune_finished_event)
+
+        await_remove = self.app._remove_nodes([self])
+        return await_remove
 
     def render(self) -> RenderableType:
         """Get renderable for widget.