await screens

CHANGELOG.md

@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## Unreleased
+
+### Changed
+
+- Breaking change: `App.push_screen` now returns an Awaitable rather than a screen. https://github.com/Textualize/textual/pull/4672
+- Breaking change: `Screen.dismiss` now returns an Awaitable rather than a bool. https://github.com/Textualize/textual/pull/4672
+
+
 ## [0.70.0] - 2024-06-19
 
 ### Fixed

src/textual/_worker_manager.py

@@ -175,5 +175,7 @@ async def wait_for_complete(self, workers: Iterable[Worker] | None = None) -> No
         Args:
             workers: An iterable of workers or None to wait for all workers in the manager.
         """
-
-        await asyncio.gather(*[worker.wait() for worker in (workers or self)])
+        try:
+            await asyncio.gather(*[worker.wait() for worker in (workers or self)])
+        except asyncio.CancelledError:
+            pass

src/textual/app.py

@@ -79,6 +79,7 @@
 from ._wait import wait_for_idle
 from ._worker_manager import WorkerManager
 from .actions import ActionParseResult, SkipAction
+from .await_complete import AwaitComplete
 from .await_remove import AwaitRemove
 from .binding import Binding, BindingType, _Bindings
 from .command import CommandPalette, Provider
@@ -1569,8 +1570,10 @@ async def run_auto_pilot(
                 if auto_pilot_task is not None:
                     await auto_pilot_task
             finally:
-                await app._shutdown()
-
+                try:
+                    await asyncio.shield(app._shutdown())
+                except asyncio.CancelledError:
+                    pass
         return app.return_value
 
     def run(
@@ -1910,7 +1913,7 @@ def add_mode(
 
         self.MODES[mode] = base_screen
 
-    def remove_mode(self, mode: str) -> None:
+    def remove_mode(self, mode: str) -> AwaitComplete:
         """Removes a mode from the app.
 
         Screens that are running in the stack of that mode are scheduled for pruning.
@@ -1930,12 +1933,17 @@ def remove_mode(self, mode: str) -> None:
             del self.MODES[mode]
 
         if mode not in self._screen_stacks:
-            return
+            return AwaitComplete.nothing()
 
         stack = self._screen_stacks[mode]
         del self._screen_stacks[mode]
-        for screen in reversed(stack):
-            self._replace_screen(screen)
+
+        async def remove_screens() -> None:
+            """Remove screens."""
+            for screen in reversed(stack):
+                await self._replace_screen(screen)
+
+        return AwaitComplete(remove_screens()).call_next(self)
 
     def is_screen_installed(self, screen: Screen | str) -> bool:
         """Check if a given screen has been installed.
@@ -2030,7 +2038,7 @@ def _load_screen_css(self, screen: Screen):
             self.stylesheet.reparse()
             self.stylesheet.update(self)
 
-    def _replace_screen(self, screen: Screen) -> Screen:
+    async def _replace_screen(self, screen: Screen) -> Screen:
         """Handle the replaced screen.
 
         Args:
@@ -2046,7 +2054,7 @@ def _replace_screen(self, screen: Screen) -> Screen:
         if not self.is_screen_installed(screen) and all(
             screen not in stack for stack in self._screen_stacks.values()
         ):
-            screen.remove()
+            await screen.remove()
             self.log.system(f"{screen} REMOVED")
         return screen
 
@@ -2151,9 +2159,10 @@ async def push_screen_wait(
         Returns:
             The screen's result.
         """
+        await self._flush_next_callbacks()
         return await self.push_screen(screen, wait_for_dismiss=True)
 
-    def switch_screen(self, screen: Screen | str) -> AwaitMount:
+    def switch_screen(self, screen: Screen | str) -> AwaitComplete:
         """Switch to another [screen](/guide/screens) by replacing the top of the screen stack with a new screen.
 
         Args:
@@ -2167,16 +2176,24 @@ def switch_screen(self, screen: Screen | str) -> AwaitMount:
         next_screen, await_mount = self._get_screen(screen)
         if screen is self.screen or next_screen is self.screen:
             self.log.system(f"Screen {screen} is already current.")
-            return AwaitMount(self.screen, [])
+            return AwaitComplete.nothing()
 
-        previous_screen = self._replace_screen(self._screen_stack.pop())
-        previous_screen._pop_result_callback()
+        top_screen = self._screen_stack.pop()
+
+        top_screen._pop_result_callback()
         self._load_screen_css(next_screen)
         self._screen_stack.append(next_screen)
         self.screen.post_message(events.ScreenResume())
         self.screen._push_result_callback(self.screen, None)
         self.log.system(f"{self.screen} is current (SWITCHED)")
-        return await_mount
+
+        async def do_switch() -> None:
+            """Task to perform switch."""
+
+            await await_mount()
+            await self._replace_screen(top_screen)
+
+        return AwaitComplete(do_switch()).call_next(self)
 
     def install_screen(self, screen: Screen, name: str) -> None:
         """Install a screen.
@@ -2238,22 +2255,29 @@ def uninstall_screen(self, screen: Screen | str) -> str | None:
                     return name
         return None
 
-    def pop_screen(self) -> Screen[object]:
+    def pop_screen(self) -> AwaitComplete:
         """Pop the current [screen](/guide/screens) from the stack, and switch to the previous screen.
 
         Returns:
             The screen that was replaced.
         """
+
         screen_stack = self._screen_stack
         if len(screen_stack) <= 1:
             raise ScreenStackError(
                 "Can't pop screen; there must be at least one screen on the stack"
             )
-        previous_screen = self._replace_screen(screen_stack.pop())
+
+        previous_screen = screen_stack.pop()
         previous_screen._pop_result_callback()
         self.screen.post_message(events.ScreenResume())
         self.log.system(f"{self.screen} is active")
-        return previous_screen
+
+        async def do_pop() -> None:
+            """Task to pop the screen."""
+            await self._replace_screen(previous_screen)
+
+        return AwaitComplete(do_pop()).call_next(self)
 
     def set_focus(self, widget: Widget | None, scroll_visible: bool = True) -> None:
         """Focus (or unfocus) a widget. A focused widget will receive key events first.
@@ -3307,12 +3331,18 @@ def _detach_from_dom(self, widgets: list[Widget]) -> list[Widget]:
         # remove and ensure that, if one of them is the focused widget,
         # focus gets moved to somewhere else.
         dedupe_to_remove = set(everything_to_remove)
-        if self.screen.focused in dedupe_to_remove:
-            self.screen._reset_focus(
-                self.screen.focused,
-                [to_remove for to_remove in dedupe_to_remove if to_remove.can_focus],
-            )
-
+        try:
+            if self.screen.focused in dedupe_to_remove:
+                self.screen._reset_focus(
+                    self.screen.focused,
+                    [
+                        to_remove
+                        for to_remove in dedupe_to_remove
+                        if to_remove.can_focus
+                    ],
+                )
+        except ScreenStackError:
+            pass
         # Next, we go through the set of widgets we've been asked to remove
         # and try and find the minimal collection of widgets that will
         # result in everything else that should be removed, being removed.
@@ -3408,7 +3438,7 @@ async def _prune_nodes(self, widgets: list[Widget]) -> None:
         """
         async with self._dom_lock:
             for widget in widgets:
-                await self._prune_node(widget)
+                await asyncio.shield(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/await_complete.py

@@ -4,6 +4,9 @@
 from typing import Any, Awaitable, Generator
 
 import rich.repr
+from typing_extensions import Self
+
+from .message_pump import MessagePump
 
 
 @rich.repr.auto(angular=True)
@@ -18,6 +21,15 @@ def __init__(self, *awaitables: Awaitable) -> None:
         """
         self._future: Future[Any] = gather(*awaitables)
 
+    def call_next(self, node: MessagePump) -> Self:
+        """Await after the next message.
+
+        Args:
+            node: The node which created the object.
+        """
+        node.call_next(self)
+        return self
+
     async def __call__(self) -> Any:
         return await self
 

src/textual/message_pump.py

@@ -493,7 +493,10 @@ async def _close_messages(self, wait: bool = True) -> None:
                 running_widget = None
 
             if running_widget is None or running_widget is not self:
-                await self._task
+                try:
+                    await self._task
+                except CancelledError:
+                    pass
 
     def _start_messages(self) -> None:
         """Start messages task."""

src/textual/screen.py

@@ -33,6 +33,7 @@
 from ._context import active_message_pump, visible_screen_stack
 from ._path import CSSPathType, _css_path_type_as_list, _make_path_object_relative
 from ._types import CallbackType
+from .await_complete import AwaitComplete
 from .binding import ActiveBinding, Binding, _Bindings
 from .css.match import match
 from .css.parse import parse_selectors
@@ -1226,43 +1227,46 @@ def _forward_event(self, event: events.Event) -> None:
     class _NoResult:
         """Class used to mark that there is no result."""
 
-    def dismiss(self, result: ScreenResultType | Type[_NoResult] = _NoResult) -> bool:
+    def dismiss(
+        self, result: ScreenResultType | Type[_NoResult] = _NoResult
+    ) -> AwaitComplete:
         """Dismiss the screen, optionally with a result.
 
         !!! note
 
             Only the active screen may be dismissed. If you try to dismiss a screen that isn't active,
-            this method will return `False`.
+            this method will raise a `ScreenError`.
 
         If `result` is provided and a callback was set when the screen was [pushed][textual.app.App.push_screen], then
         the callback will be invoked with `result`.
 
         Args:
             result: The optional result to be passed to the result callback.
 
-        Returns:
-            `True` if the Screen was dismissed, or `False` if the Screen wasn't dismissed due to not being active.
-
         Raises:
+            ScreenError: If the screen being dismissed is not active.
             ScreenStackError: If trying to dismiss a screen that is not at the top of
                 the stack.
 
         """
         if not self.is_active:
-            return False
+            from .app import ScreenError
+
+            raise ScreenError("Screen is not active")
         if result is not self._NoResult and self._result_callbacks:
             self._result_callbacks[-1](cast(ScreenResultType, result))
-        self.app.pop_screen()
-        return True
+        await_pop = self.app.pop_screen()
+        return await_pop
 
-    def action_dismiss(
+    async def action_dismiss(
         self, result: ScreenResultType | Type[_NoResult] = _NoResult
     ) -> None:
         """A wrapper around [`dismiss`][textual.screen.Screen.dismiss] that can be called as an action.
 
         Args:
             result: The optional result to be passed to the result callback.
         """
+        await self._flush_next_callbacks()
         self.dismiss(result)
 
     def can_view(self, widget: Widget) -> bool:

src/textual/widget.py

@@ -676,7 +676,7 @@ def set_loading(self, loading: bool) -> Awaitable:
             loading_indicator = self.get_loading_widget()
             loading_indicator.add_class(LOADING_INDICATOR_CLASS)
             await_mount = self.mount(loading_indicator)
-            return AwaitComplete(remove_indicator, await_mount)
+            return AwaitComplete(remove_indicator, await_mount).call_next(self)
         else:
             return remove_indicator
 

tests/css/test_screen_css.py

@@ -210,7 +210,7 @@ async def test_screen_css_switch_screen_type_by_name():
     class MyApp(SwitchBaseApp):
         SCREENS = {"screenwithcss": ScreenWithCSS}
 
-        def key_p(self):
+        async def key_p(self):
             self.switch_screen("screenwithcss")
 
         def key_o(self):

tests/test_screens.py

@@ -5,7 +5,7 @@
 import pytest
 
 from textual import work
-from textual.app import App, ComposeResult, ScreenStackError
+from textual.app import App, ComposeResult, ScreenError, ScreenStackError
 from textual.events import MouseMove
 from textual.geometry import Offset
 from textual.screen import Screen
@@ -110,7 +110,7 @@ async def test_screens():
     # Check screen stack is empty
     assert app.screen_stack == []
     # Push a screen
-    app.push_screen("screen1")
+    await app.push_screen("screen1")
     # Check it is on the stack
     assert app.screen_stack == [screen1]
     # Check it is current
@@ -119,21 +119,21 @@ async def test_screens():
     assert app.children == (screen1,)
 
     # Switch to another screen
-    app.switch_screen("screen2")
+    await app.switch_screen("screen2")
     # Check it has changed the stack and that it is current
     assert app.screen_stack == [screen2]
     assert app.screen is screen2
     assert app.children == (screen2,)
 
     # Push another screen
-    app.push_screen("screen3")
+    await app.push_screen("screen3")
     assert app.screen_stack == [screen2, screen3]
     assert app.screen is screen3
     # Only the current screen is in children
     assert app.children == (screen3,)
 
     # Pop a screen
-    assert app.pop_screen() is screen3
+    await app.pop_screen()
     assert app.screen is screen2
     assert app.screen_stack == [screen2]
 
@@ -293,14 +293,16 @@ def on_mount(self):
 async def test_dismiss_non_top_screen():
     class MyApp(App[None]):
         async def key_p(self) -> None:
-            self.bottom, top = Screen(), Screen()
+            self.bottom = Screen()
+            top = Screen()
             await self.push_screen(self.bottom)
             await self.push_screen(top)
 
     app = MyApp()
     async with app.run_test() as pilot:
         await pilot.press("p")
-        assert not app.bottom.dismiss()
+        with pytest.raises(ScreenError):
+            await app.bottom.dismiss()
 
 
 async def test_dismiss_action():