allow waiting of screen operations

Author: willmcgugan

src/textual/app.py

@@ -860,8 +860,6 @@ def is_screen_installed(self, screen: Screen | str) -> bool:
     def get_screen(self, screen: Screen | str) -> Screen:
         """Get an installed screen.
 
-        If the screen isn't running, it will be registered before it is run.
-
         Args:
             screen (Screen | str): Either a Screen object or screen name (the `name` argument when installed).
 
@@ -878,10 +876,30 @@ def get_screen(self, screen: Screen | str) -> Screen:
                 raise KeyError(f"No screen called {screen!r} installed") from None
         else:
             next_screen = screen
-        if not next_screen.is_running:
-            self._register(self, next_screen)
         return next_screen
 
+    def _get_screen(self, screen: Screen | str) -> tuple[Screen, AwaitMount]:
+        """Get an installed screen and a await mount object.
+
+        If the screen isn't running, it will be registered before it is run.
+
+        Args:
+            screen (Screen | str): Either a Screen object or screen name (the `name` argument when installed).
+
+        Raises:
+            KeyError: If the named screen doesn't exist.
+
+        Returns:
+            tuple[Screen, AwaitMount]: A screen instance and an awaitable that awaits the children mounting.
+
+        """
+        _screen = self.get_screen(screen)
+        if not _screen.is_running:
+            widgets = self._register(self, _screen)
+            return (_screen, AwaitMount(widgets))
+        else:
+            return (_screen, AwaitMount([]))
+
     def _replace_screen(self, screen: Screen) -> Screen:
         """Handle the replaced screen.
 
@@ -899,19 +917,20 @@ def _replace_screen(self, screen: Screen) -> Screen:
             self.log.system(f"{screen} REMOVED")
         return screen
 
-    def push_screen(self, screen: Screen | str) -> None:
+    def push_screen(self, screen: Screen | str) -> AwaitMount:
         """Push a new screen on the screen stack.
 
         Args:
             screen (Screen | str): A Screen instance or the name of an installed screen.
 
         """
-        next_screen = self.get_screen(screen)
+        next_screen, await_mount = self._get_screen(screen)
         self._screen_stack.append(next_screen)
         self.screen.post_message_no_wait(events.ScreenResume(self))
         self.log.system(f"{self.screen} is current (PUSHED)")
+        return await_mount
 
-    def switch_screen(self, screen: Screen | str) -> None:
+    def switch_screen(self, screen: Screen | str) -> AwaitMount:
         """Switch to another screen by replacing the top of the screen stack with a new screen.
 
         Args:
@@ -920,12 +939,14 @@ def switch_screen(self, screen: Screen | str) -> None:
         """
         if self.screen is not screen:
             self._replace_screen(self._screen_stack.pop())
-            next_screen = self.get_screen(screen)
+            next_screen, await_mount = self._get_screen(screen)
             self._screen_stack.append(next_screen)
             self.screen.post_message_no_wait(events.ScreenResume(self))
             self.log.system(f"{self.screen} is current (SWITCHED)")
+            return await_mount
+        return AwaitMount([])
 
-    def install_screen(self, screen: Screen, name: str | None = None) -> str:
+    def install_screen(self, screen: Screen, name: str | None = None) -> AwaitMount:
         """Install a screen.
 
         Args:
@@ -937,7 +958,7 @@ def install_screen(self, screen: Screen, name: str | None = None) -> str:
             ScreenError: If the screen can't be installed.
 
         Returns:
-            str: The name of the screen
+            AwaitMount: An awaitable that awaits the mounting of the screen and its children.
         """
         if name is None:
             name = nanoid.generate()
@@ -948,9 +969,9 @@ def install_screen(self, screen: Screen, name: str | None = None) -> str:
                 "Can't install screen; {screen!r} has already been installed"
             )
         self._installed_screens[name] = screen
-        self.get_screen(name)  # Ensures screen is running
+        _screen, await_mount = self._get_screen(name)  # Ensures screen is running
         self.log.system(f"{screen} INSTALLED name={name!r}")
-        return name
+        return await_mount
 
     def uninstall_screen(self, screen: Screen | str) -> str | None:
         """Uninstall a screen. If the screen was not previously installed then this

src/textual/demo.css

@@ -114,13 +114,14 @@ DarkSwitch {
 }
 
 DarkSwitch .label {
-    
+    width: 1fr;
     padding: 1 2;
     color: $text-muted;
 }
 
 DarkSwitch Checkbox {
     background: $boost;
+    dock: left;
 }
 
 

src/textual/screen.py

@@ -12,7 +12,6 @@
 from ._compositor import Compositor, MapGeometry
 from .timer import Timer
 from ._types import CallbackType
-from .dom import DOMNode
 from .geometry import Offset, Region, Size
 from .reactive import Reactive
 from .renderables.blank import Blank
@@ -61,7 +60,12 @@ def is_transparent(self) -> bool:
     @property
     def is_current(self) -> bool:
         """Check if this screen is current (i.e. visible to user)."""
-        return self.app.screen is self
+        from .app import ScreenStackError
+
+        try:
+            return self.app.screen is self
+        except ScreenStackError:
+            return False
 
     @property
     def update_timer(self) -> Timer:

src/textual/widget.py

@@ -82,11 +82,13 @@ def __init__(self, widgets: Sequence[Widget]) -> None:
 
     def __await__(self) -> Generator[None, None, None]:
         async def await_mount() -> None:
-            aws = [
-                create_task(widget._mounted_event.wait()) for widget in self._widgets
-            ]
-            if aws:
-                await wait(aws)
+            if self._widgets:
+                aws = [
+                    create_task(widget._mounted_event.wait())
+                    for widget in self._widgets
+                ]
+                if aws:
+                    await wait(aws)
 
         return await_mount().__await__()
 

tests/test_unmount.py

@@ -29,11 +29,10 @@ def on_unmount(self, event: events.Unmount):
 
     class UnmountApp(App):
         async def on_mount(self) -> None:
-            self.push_screen(MyScreen(id="main"))
+            await self.push_screen(MyScreen(id="main"))
 
     app = UnmountApp()
     async with app.run_test() as pilot:
-        await pilot.pause()  # TODO remove when push_screen is awaitable
         await pilot.exit(None)
 
     expected = [