Stop mounting children from blocking the parent's message queue

src/textual/app.py

@@ -2822,11 +2822,11 @@ def _get_screen(self, screen: Screen | str) -> tuple[Screen, AwaitMount]:
         if not _screen.is_running:
             widgets = self._register(self, _screen)
             await_mount = AwaitMount(_screen, widgets)
-            self.call_next(await_mount)
+            await_mount.call_when_ready(self)
             return (_screen, await_mount)
         else:
             await_mount = AwaitMount(_screen, [])
-            self.call_next(await_mount)
+            await_mount.call_when_ready(self)
             return (_screen, await_mount)
 
     def _load_screen_css(self, screen: Screen):

src/textual/await_complete.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from asyncio import Future, gather
 from typing import TYPE_CHECKING, Any, Awaitable, Generator
 
@@ -13,27 +14,18 @@
     from textual.types import CallbackType
 
 
-@rich.repr.auto(angular=True)
-class AwaitComplete:
-    """An 'optionally-awaitable' object which runs one or more coroutines (or other awaitables) concurrently."""
+class AwaitCompletion(ABC):
+    """Shared receipt behavior, independent of how its operation is acquired.
 
-    def __init__(
-        self, *awaitables: Awaitable, pre_await: CallbackType | None = None
-    ) -> None:
-        """Create an AwaitComplete.
+    Coroutine groups, mounts and removals have distinct constructors and wait
+    lifetimes. They share completion notification and pre-await admission,
+    not a zero-argument operation factory or cancellation policy.
+    """
 
-        Args:
-            awaitables: One or more awaitables to run concurrently.
-        """
-        self._awaitables = awaitables
-        self._future: Future[Any] = gather(*awaitables)
+    def __init__(self, *, pre_await: CallbackType | None = None) -> None:
         self._pre_await: CallbackType | None = pre_await
         self._caller = get_caller_file_and_line()
-
-    def __rich_repr__(self) -> rich.repr.Result:
-        yield self._awaitables
-        yield "pre_await", self._pre_await, None
-        yield "caller", self._caller, None
+        self._scheduled = False
 
     def set_pre_await_callback(self, pre_await: CallbackType | None) -> None:
         """Set a callback to run prior to awaiting.
@@ -55,27 +47,79 @@ def call_next(self, node: MessagePump) -> Self:
         node.call_next(self)
         return self
 
+    @abstractmethod
+    def _start(self) -> Future[Any]:
+        """The completion owned by this optional awaitable."""
+
+    @abstractmethod
+    def _await(self) -> Awaitable[Any]:
+        """Acquire the concrete operation's wait / cancellation contract."""
+
+    def call_when_ready(self, node: MessagePump) -> None:
+        """Deliver completion without occupying the receiver's message pump.
+
+        Mount and removal retain their own completion lifetimes. The original
+        callback path observes their result only when awaiting it cannot block
+        input or other messages behind another widget's startup / teardown.
+        """
+        if self._scheduled:
+            return
+        self._scheduled = True
+
+        def completed(future: Future[Any]) -> None:
+            if node._closing or node._closed:
+                if not future.cancelled():
+                    future.exception()
+                return
+            node.call_next(self)
+
+        self._start().add_done_callback(completed)
+
     async def __call__(self) -> Any:
         return await self
 
     def __await__(self) -> Generator[Any, None, Any]:
         _rich_traceback_omit = True
         if self._pre_await is not None:
             self._pre_await()
-        return self._future.__await__()
+        return self._await().__await__()
 
     @property
     def is_done(self) -> bool:
         """`True` if the task has completed."""
-        return self._future.done()
+        return self._start().done()
 
     @property
     def exception(self) -> BaseException | None:
         """An exception if the awaitables failed."""
-        if self._future.done():
-            return self._future.exception()
+        completion = self._start()
+        if completion.done():
+            return completion.exception()
         return None
 
+
+@rich.repr.auto(angular=True)
+class AwaitComplete(AwaitCompletion):
+    """An optional awaitable which runs a group of awaitables concurrently."""
+
+    def __init__(
+        self, *awaitables: Awaitable, pre_await: CallbackType | None = None
+    ) -> None:
+        super().__init__(pre_await=pre_await)
+        self._awaitables = awaitables
+        self._future: Future[Any] = gather(*awaitables)
+
+    def __rich_repr__(self) -> rich.repr.Result:
+        yield self._awaitables
+        yield "pre_await", self._pre_await, None
+        yield "caller", self._caller, None
+
+    def _start(self) -> Future[Any]:
+        return self._future
+
+    def _await(self) -> Awaitable[Any]:
+        return self._future
+
     @classmethod
     def nothing(cls):
         """Returns an already completed instance of AwaitComplete."""

src/textual/await_remove.py

@@ -6,34 +6,32 @@
 
 import asyncio
 from asyncio import Future, Task, gather
-from typing import TYPE_CHECKING, Generator
+from typing import TYPE_CHECKING, Awaitable
 
 import rich.repr
 
 from textual._callback import invoke
-from textual._debug import get_caller_file_and_line
 from textual._types import CallbackType
+from textual.await_complete import AwaitCompletion
 
 if TYPE_CHECKING:
     from textual.dom import DOMNode
-    from textual.message_pump import MessagePump
     from textual.screen import Screen
     from textual.widget import Widget
 
 
 @rich.repr.auto
-class AwaitRemove:
+class AwaitRemove(AwaitCompletion):
     """An awaitable that waits for nodes to be removed."""
 
     def __init__(
         self, tasks: list[Task], post_remove: CallbackType | None = None
     ) -> None:
+        super().__init__()
         self._tasks = list(tasks)
         self._post_remove = post_remove
-        self._caller = get_caller_file_and_line()
         self._completion: Future[None] | None = None
         self._finisher: Task[None] | None = None
-        self._scheduled = False
 
     @classmethod
     def prune(cls, *nodes: Widget, parent: DOMNode | None = None) -> AwaitRemove:
@@ -92,9 +90,6 @@ def __rich_repr__(self) -> rich.repr.Result:
         yield "post_remove", self._post_remove
         yield "caller", self._caller, None
 
-    async def __call__(self) -> None:
-        await self
-
     def _start(self) -> Future[None]:
         """One independently-owned teardown completion for all optional waiters."""
         if self._completion is None:
@@ -121,27 +116,7 @@ async def _finish(self) -> None:
             self._post_remove = None
             self._finisher = None
 
-    def call_when_ready(self, node: MessagePump) -> None:
-        """Deliver the completed receipt without blocking the receiver's queue.
-
-        The application's message pump must remain available while another
-        widget runs an asynchronous unmount handler. Observe errors through the
-        ordinary callback path only after this transaction has actually ended.
-        """
-        if self._scheduled:
-            return
-        self._scheduled = True
-
-        def completed(future: Future[None]) -> None:
-            if node._closing or node._closed:
-                if not future.cancelled():
-                    future.exception()
-                return
-            node.call_next(self)
-
-        self._start().add_done_callback(completed)
-
-    def __await__(self) -> Generator[None, None, None]:
+    def _await(self) -> Awaitable[None]:
         current_task = asyncio.current_task()
         self_removal = current_task in self._tasks
         other_tasks = [task for task in self._tasks if task is not current_task]
@@ -156,4 +131,4 @@ async def await_prune() -> None:
                 return
             await asyncio.shield(completion)
 
-        return await_prune().__await__()
+        return await_prune()

src/textual/widget.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from asyncio import Lock, create_task, gather, wait
+from asyncio import Task, create_task, gather, shield
 from collections import Counter
 from contextlib import asynccontextmanager
 from fractions import Fraction
@@ -14,10 +14,10 @@
 from typing import (
     TYPE_CHECKING,
     AsyncGenerator,
+    Awaitable,
     Callable,
     ClassVar,
     Collection,
-    Generator,
     Iterable,
     Mapping,
     NamedTuple,
@@ -51,7 +51,6 @@
 from textual._animator import DEFAULT_EASING, Animatable, BoundAnimator, EasingFunction
 from textual._arrange import DockArrangeResult, arrange
 from textual._context import NoActiveAppError
-from textual._debug import get_caller_file_and_line
 from textual._dispatch_key import dispatch_key
 from textual._easing import DEFAULT_SCROLL_EASING
 from textual._extrema import Extrema
@@ -65,7 +64,7 @@
 from textual._styles_cache import StylesCache
 from textual._types import AnimationLevel
 from textual.actions import SkipAction
-from textual.await_complete import AwaitComplete
+from textual.await_complete import AwaitComplete, AwaitCompletion
 from textual.await_remove import AwaitRemove
 from textual.box_model import BoxModel
 from textual.cache import FIFOCache, LRUCache
@@ -132,7 +131,7 @@
 
 
 @rich.repr.auto
-class AwaitMount:
+class AwaitMount(AwaitCompletion):
     """An *optional* awaitable returned by [mount][textual.widget.Widget.mount] and [mount_all][textual.widget.Widget.mount_all].
 
     Example:
@@ -142,50 +141,34 @@ class AwaitMount:
     """
 
     def __init__(self, parent: Widget, widgets: Sequence[Widget]) -> None:
+        super().__init__()
         self._parent = parent
         self._widgets = widgets
-        self._caller = get_caller_file_and_line()
-        self._completion_lock = Lock()
-        self._completed = False
+        self._future: Task[None] = create_task(self._finish(), name="complete mount")
 
     def __rich_repr__(self) -> rich.repr.Result:
         yield "parent", self._parent
         yield "widgets", self._widgets
         yield "caller", self._caller, None
 
-    async def __call__(self) -> None:
-        """Allows awaiting via a call operation."""
-        await self
-
-    def __await__(self) -> Generator[None, None, None]:
-        async def await_mount() -> None:
-            # mount() schedules this same object with call_next, and the caller
-            # may await it explicitly too. Complete one mount transaction, not
-            # a second set of waiters and layout invalidations for each await.
-            async with self._completion_lock:
-                if self._completed:
-                    return
-                if self._widgets:
-                    aws = [
-                        create_task(widget._mounted_event.wait(), name="await mount")
-                        for widget in self._widgets
-                    ]
-                    try:
-                        await wait(aws)
-                    finally:
-                        pending = [task for task in aws if not task.done()]
-                        if pending:
-                            for task in pending:
-                                task.cancel()
-                            await gather(*pending, return_exceptions=True)
-                    self._parent.refresh(layout=True)
-                    try:
-                        self._parent.app._update_mouse_over(self._parent.screen)
-                    except NoScreen:
-                        pass
-                self._completed = True
-
-        return await_mount().__await__()
+    async def _finish(self) -> None:
+        """Complete registered child startup once, independently of waiters."""
+        if self._widgets:
+            await gather(*(widget._mounted_event.wait() for widget in self._widgets))
+            if not (self._parent._closing or self._parent._closed or self._parent._pruning):
+                self._parent.refresh(layout=True)
+                try:
+                    self._parent.app._update_mouse_over(self._parent.screen)
+                except NoScreen:
+                    pass
+
+    def _start(self) -> Task[None]:
+        return self._future
+
+    def _await(self) -> Awaitable[None]:
+        # Cancelling an optional caller cannot cancel registered child startup
+        # or its completion publication to other callers / the parent.
+        return shield(self._future)
 
 
 class _Styled:
@@ -1591,7 +1574,7 @@ def update_styles(children: list[DOMNode]) -> None:
 
         self.call_later(update_styles, self.displayed_children)
         await_mount = AwaitMount(self, mounted)
-        self.call_next(await_mount)
+        await_mount.call_when_ready(self)
 
         return await_mount