Fix hang when removing current widget

src/textual/_context.py

@@ -4,6 +4,7 @@ from contextvars import ContextVar
 
 if TYPE_CHECKING:
     from .app import App
+    from .message_pump import MessagePump
 
 
 class NoActiveAppError(RuntimeError):
@@ -11,3 +12,4 @@ class NoActiveAppError(RuntimeError):
 
 
 active_app: ContextVar["App"] = ContextVar("active_app")
+active_message_pump: ContextVar["MessagePump"] = ContextVar("active_message_pump")

src/textual/app.py

@@ -46,7 +46,7 @@ from ._animator import DEFAULT_EASING, Animatable, Animator, EasingFunction
 from ._ansi_sequences import SYNC_END, SYNC_START
 from ._asyncio import create_task
 from ._callback import invoke
-from ._context import active_app
+from ._context import active_app, active_message_pump
 from ._event_broker import NoHandler, extract_handler_actions
 from ._filter import LineFilter, Monochrome
 from ._path import _make_path_object_relative
@@ -1113,6 +1113,7 @@ class App(Generic[ReturnType], DOMNode):
     def mount_all(
         self,
         widgets: Iterable[Widget],
+        *,
         before: int | str | Widget | None = None,
         after: int | str | Widget | None = None,
     ) -> AwaitMount:
@@ -2103,7 +2104,7 @@ class App(Generic[ReturnType], DOMNode):
         """Remove nodes from DOM, and return an awaitable that awaits cleanup.
 
         Args:
-            widgets: List of nodes to remvoe.
+            widgets: List of nodes to remove.
 
         Returns:
             Awaitable that returns when the nodes have been fully removed.
@@ -2131,13 +2132,15 @@ class App(Generic[ReturnType], DOMNode):
             prune_widgets_task(removed_widgets, finished_event), name="prune nodes"
         )
 
-        return AwaitRemove(finished_event)
+        await_remove = AwaitRemove(finished_event)
+        self.call_next(await_remove)
+        return await_remove
 
     async def _prune_nodes(self, widgets: list[Widget]) -> None:
         """Remove nodes and children.
 
         Args:
-            widgets: _description_
+            widgets: Widgets to remove.
         """
         async with self._dom_lock:
             for widget in widgets:

src/textual/await_remove.py

@@ -15,6 +15,9 @@ class AwaitRemove:
         """
         self.finished_flag = finished_flag
 
+    async def __call__(self) -> None:
+        return await self
+
     def __await__(self) -> Generator[None, None, None]:
         async def await_prune() -> None:
             """Wait for the prune operation to finish."""

src/textual/message_pump.py

@@ -17,7 +17,7 @@ from weakref import WeakSet
 from . import Logger, events, log, messages
 from ._asyncio import create_task
 from ._callback import invoke
-from ._context import NoActiveAppError, active_app
+from ._context import NoActiveAppError, active_app, active_message_pump
 from ._time import time
 from .case import camel_to_snake
 from .errors import DuplicateKeyHandlers
@@ -313,8 +313,13 @@ class MessagePump(metaclass=MessagePumpMeta):
         Reactive._reset_object(self)
         await self._message_queue.put(None)
         if wait and self._task is not None and asyncio.current_task() != self._task:
-            # Ensure everything is closed before returning
-            await self._task
+            try:
+                running_widget = active_message_pump.get()
+            except LookupError:
+                await self._task
+            else:
+                if running_widget is not self:
+                    await self._task
 
     def _start_messages(self) -> None:
         """Start messages task."""
@@ -382,7 +387,7 @@ class MessagePump(metaclass=MessagePumpMeta):
                     break
 
             self._active_message = message
-
+            context_token = active_message_pump.set(self)
             try:
                 try:
                     await self._dispatch_message(message)
@@ -415,6 +420,7 @@ class MessagePump(metaclass=MessagePumpMeta):
                                     self.app._handle_exception(error)
                                     break
             finally:
+                active_message_pump.reset(context_token)
                 self._active_message = None
 
     async def _flush_next_callbacks(self) -> None:

src/textual/widget.py

@@ -618,6 +618,33 @@ class Widget(DOMNode):
         self.call_next(await_mount)
         return await_mount
 
+    def mount_all(
+        self,
+        widgets: Iterable[Widget],
+        *,
+        before: int | str | Widget | None = None,
+        after: int | str | Widget | None = None,
+    ) -> AwaitMount:
+        """Mount widgets from an iterable.
+
+        Args:
+            widgets: An iterable of widgets.
+            before: Optional location to mount before.
+            after: Optional location to mount after.
+
+        Returns:
+            An awaitable object that waits for widgets to be mounted.
+
+        Raises:
+            MountError: If there is a problem with the mount request.
+
+        Note:
+            Only one of ``before`` or ``after`` can be provided. If both are
+            provided a ``MountError`` will be raised.
+        """
+        await_mount = self.mount(*widgets, before=before, after=after)
+        return await_mount
+
     def move_child(
         self,
         child: int | Widget,