From 6d1a37d57ef90252923332b3feba28b074320d6b Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 17:52:07 +0100
Subject: [PATCH 01/11] docstrings

---
 src/textual/app.py    | 60 ++++++++++++++++++++++++++++++++++++++++---
 src/textual/screen.py | 11 +++++++-
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index bb06942226..424052e94a 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -99,10 +99,11 @@
 from .screen import Screen, ScreenResultCallbackType, ScreenResultType
 from .widget import AwaitMount, Widget
 from .widgets._toast import ToastRack
+from .worker import NoActiveWorker, get_current_worker
 
 if TYPE_CHECKING:
     from textual_dev.client import DevtoolsClient
-    from typing_extensions import Coroutine, TypeAlias
+    from typing_extensions import Coroutine, Literal, TypeAlias
 
     from ._types import MessageTarget
 
@@ -246,6 +247,24 @@ def fileno(self) -> int:
         return -1
 
 
+class ScreenAwaitable(Generic[ScreenResultType]):
+    """An optional awaitable to get the result of a screen."""
+
+    def __init__(
+        self, await_mount: AwaitMount, future: asyncio.Future[ScreenResultType]
+    ) -> None:
+        self._await_mount = await_mount
+        self._future = future
+
+    def __await__(self) -> Generator[None, None, ScreenResultType]:
+        async def await_screen() -> ScreenResultType:
+            await self._await_mount
+            await self._future
+            return self._future.result()
+
+        return await_screen().__await__()
+
+
 @rich.repr.auto
 class App(Generic[ReturnType], DOMNode):
     """The base class for Textual Applications."""
@@ -1778,16 +1797,37 @@ def _replace_screen(self, screen: Screen) -> Screen:
             self.log.system(f"{screen} REMOVED")
         return screen
 
+    @overload
     def push_screen(
         self,
         screen: Screen[ScreenResultType] | str,
         callback: ScreenResultCallbackType[ScreenResultType] | None = None,
+        wait_for_dismiss: Literal[False] = False,
     ) -> AwaitMount:
+        ...
+
+    @overload
+    def push_screen(
+        self,
+        screen: Screen[ScreenResultType] | str,
+        callback: ScreenResultCallbackType[ScreenResultType] | None = None,
+        wait_for_dismiss: Literal[True] = True,
+    ) -> ScreenAwaitable[ScreenResultType]:
+        ...
+
+    def push_screen(
+        self,
+        screen: Screen[ScreenResultType] | str,
+        callback: ScreenResultCallbackType[ScreenResultType] | None = None,
+        wait_for_dismiss: bool = False,
+    ) -> AwaitMount | ScreenAwaitable[ScreenResultType]:
         """Push a new [screen](/guide/screens) on the screen stack, making it the current screen.
 
         Args:
             screen: A Screen instance or the name of an installed screen.
             callback: An optional callback function that will be called if the screen is [dismissed][textual.screen.Screen.dismiss] with a result.
+            wait_for_dismiss: If `True`, awaiting this method will return the dismiss value from the screen. When set to `False`, awaiting
+                this method will wait for the screen to be mounted. Note that `wait_for_dismiss` should only be set to `True` when running in a worker.
 
         Returns:
             An optional awaitable that awaits the mounting of the screen and its children.
@@ -1797,6 +1837,9 @@ def push_screen(
                 f"push_screen requires a Screen instance or str; not {screen!r}"
             )
 
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+
         if self._screen_stack:
             self.screen.post_message(events.ScreenSuspend())
             self.screen.refresh()
@@ -1805,13 +1848,24 @@ def push_screen(
             message_pump = active_message_pump.get()
         except LookupError:
             message_pump = self.app
-        next_screen._push_result_callback(message_pump, callback)
+
+        next_screen._push_result_callback(message_pump, callback, future)
         self._load_screen_css(next_screen)
         self._screen_stack.append(next_screen)
         self.stylesheet.update(next_screen)
         next_screen.post_message(events.ScreenResume())
         self.log.system(f"{self.screen} is current (PUSHED)")
-        return await_mount
+        if wait_for_dismiss:
+            try:
+                get_current_worker()
+            except NoActiveWorker:
+                raise NoActiveWorker(
+                    "push_screen must be run from a worker when `wait_for_dismiss` is True"
+                ) from None
+            screen_awaitable = ScreenAwaitable(await_mount, future)
+            return screen_awaitable
+        else:
+            return await_mount
 
     def switch_screen(self, screen: Screen | str) -> AwaitMount:
         """Switch to another [screen](/guide/screens) by replacing the top of the screen stack with a new screen.
diff --git a/src/textual/screen.py b/src/textual/screen.py
index c5404bc05f..89f26f504e 100644
--- a/src/textual/screen.py
+++ b/src/textual/screen.py
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from functools import partial
 from operator import attrgetter
 from typing import (
@@ -74,17 +75,21 @@ def __init__(
         self,
         requester: MessagePump,
         callback: ScreenResultCallbackType[ScreenResultType] | None,
+        future: asyncio.Future[ScreenResultType] | None = None,
     ) -> None:
         """Initialise the result callback object.
 
         Args:
             requester: The object making a request for the callback.
             callback: The callback function.
+            future: A Future to hold the result
         """
         self.requester = requester
         """The object in the DOM that requested the callback."""
         self.callback: ScreenResultCallbackType | None = callback
         """The callback function."""
+        self.future = future
+        """A future for the result"""
 
     def __call__(self, result: ScreenResultType) -> None:
         """Call the callback, passing the given result.
@@ -95,6 +100,8 @@ def __call__(self, result: ScreenResultType) -> None:
         Note:
             If the requested or the callback are `None` this will be a no-op.
         """
+        if self.future is not None:
+            self.future.set_result(result)
         if self.requester is not None and self.callback is not None:
             self.requester.call_next(self.callback, result)
 
@@ -687,15 +694,17 @@ def _push_result_callback(
         self,
         requester: MessagePump,
         callback: ScreenResultCallbackType[ScreenResultType] | None,
+        future: asyncio.Future[ScreenResultType] | None = None,
     ) -> None:
         """Add a result callback to the screen.
 
         Args:
             requester: The object requesting the callback.
             callback: The callback.
+            future: A Future to hold the result.
         """
         self._result_callbacks.append(
-            ResultCallback[ScreenResultType](requester, callback)
+            ResultCallback[ScreenResultType](requester, callback, future)
         )
 
     def _pop_result_callback(self) -> None:

From 7cc84b9a398f7fe3b98a8c6cd7d195c9daa9d534 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 17:59:52 +0100
Subject: [PATCH 02/11] raises docstring

---
 src/textual/app.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/textual/app.py b/src/textual/app.py
index 424052e94a..b5bb4ffe3f 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -1829,6 +1829,9 @@ def push_screen(
             wait_for_dismiss: If `True`, awaiting this method will return the dismiss value from the screen. When set to `False`, awaiting
                 this method will wait for the screen to be mounted. Note that `wait_for_dismiss` should only be set to `True` when running in a worker.
 
+        Raises:
+            NoActiveWorker: If using `wait_for_dismiss` outside of a worker.
+
         Returns:
             An optional awaitable that awaits the mounting of the screen and its children.
         """

From 154e67c439127ff434d072a6cb60860b018e6a99 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 18:12:48 +0100
Subject: [PATCH 03/11] fix for tests

---
 src/textual/app.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index b5bb4ffe3f..480d88b349 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -1840,8 +1840,13 @@ def push_screen(
                 f"push_screen requires a Screen instance or str; not {screen!r}"
             )
 
-        loop = asyncio.get_running_loop()
-        future = loop.create_future()
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            # Mainly for testing, when push_screen isn't called in an async context
+            future: asyncio.Future[ScreenResultType] = asyncio.Future()
+        else:
+            future = loop.create_future()
 
         if self._screen_stack:
             self.screen.post_message(events.ScreenSuspend())

From 5c9fa69f0e6e4fbcca8f778f726597ba2ead1dd4 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 18:27:55 +0100
Subject: [PATCH 04/11] Formatting

---
 src/textual/app.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index 480d88b349..4a532fac38 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -274,17 +274,14 @@ class App(Generic[ReturnType], DOMNode):
     and therefore takes priority in the event of a specificity clash."""
 
     # Default (the lowest priority) CSS
-    DEFAULT_CSS: ClassVar[
-        str
-    ] = """
+    DEFAULT_CSS: ClassVar[str]
+    DEFAULT_CSS = """
     App {
         background: $background;
         color: $text;
     }
-
     *:disabled:can-focus {
         opacity: 0.7;
-
     }
     """
 

From fb7f24b21205cae820baeccb937ca4ae60c5ef90 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 19:36:22 +0100
Subject: [PATCH 05/11] tests

---
 src/textual/app.py    |  1 +
 tests/test_screens.py | 72 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index 4a532fac38..288b1beaed 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -258,6 +258,7 @@ def __init__(
 
     def __await__(self) -> Generator[None, None, ScreenResultType]:
         async def await_screen() -> ScreenResultType:
+            """Await the mount and the future."""
             await self._await_mount
             await self._future
             return self._future.result()
diff --git a/tests/test_screens.py b/tests/test_screens.py
index 5f587fd0ee..aa2a3c80a1 100644
--- a/tests/test_screens.py
+++ b/tests/test_screens.py
@@ -4,11 +4,13 @@
 
 import pytest
 
-from textual.app import App, ScreenStackError, ComposeResult
+from textual import work
+from textual.app import App, ComposeResult, ScreenStackError
 from textual.events import MouseMove
 from textual.geometry import Offset
 from textual.screen import Screen
 from textual.widgets import Button, Input, Label
+from textual.worker import NoActiveWorker
 
 skip_py310 = pytest.mark.skipif(
     sys.version_info.minor == 10 and sys.version_info.major == 3,
@@ -407,4 +409,70 @@ def on_mount(self):
 
     assert len(MouseMoveRecordingScreen.mouse_events) == 1
     mouse_event = MouseMoveRecordingScreen.mouse_events[0]
-    assert mouse_event.x, mouse_event.y == (label_offset.x + mouse_offset.x, label_offset.y + mouse_offset.y)
+    assert mouse_event.x, mouse_event.y == (
+        label_offset.x + mouse_offset.x,
+        label_offset.y + mouse_offset.y,
+    )
+
+
+async def test_push_screen_wait_for_dismiss() -> None:
+    """Test push_screen returns result."""
+
+    class QuitScreen(Screen):
+        BINDINGS = [
+            ("y", "quit(True)"),
+            ("n", "quit(False)"),
+        ]
+
+        def action_quit(self, quit: bool) -> None:
+            self.dismiss(quit)
+
+    results: list[bool] = []
+
+    class ScreensApp(App):
+        BINDINGS = [("x", "exit")]
+
+        @work
+        async def action_exit(self) -> None:
+            result = await self.push_screen(QuitScreen(), wait_for_dismiss=True)
+            results.append(result)
+
+    app = ScreensApp()
+    # Press X to exit, then Y to dismiss, expect True result
+    async with app.run_test() as pilot:
+        await pilot.press("x", "y")
+    assert results == [True]
+
+    results.clear()
+    app = ScreensApp()
+    # Press X to exit, then Y to dismiss, expect True result
+    async with app.run_test() as pilot:
+        await pilot.press("x", "n")
+    assert results == [False]
+
+
+async def test_push_screen_wait_for_dismiss_no_worker() -> None:
+    """Test wait_for_dismiss raises NoActiveWorker when not using workers."""
+
+    class QuitScreen(Screen):
+        BINDINGS = [
+            ("y", "quit(True)"),
+            ("n", "quit(False)"),
+        ]
+
+        def action_quit(self, quit: bool) -> None:
+            self.dismiss(quit)
+
+    results: list[bool] = []
+
+    class ScreensApp(App):
+        BINDINGS = [("x", "exit")]
+
+        async def action_exit(self) -> None:
+            result = await self.push_screen(QuitScreen(), wait_for_dismiss=True)
+            results.append(result)
+
+    app = ScreensApp()
+    with pytest.raises(NoActiveWorker):
+        async with app.run_test() as pilot:
+            await pilot.press("x", "y")

From c5cffb8468478444e9a191ffaf8aae98683ad435 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 19:37:09 +0100
Subject: [PATCH 06/11] changelog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4c6825debc..23f150e699 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Reactive `cell_padding` (and respective parameter) to define horizontal cell padding in data table columns https://github.com/Textualize/textual/issues/3435
 - Added `Input.clear` method https://github.com/Textualize/textual/pull/3430
 - Added `TextArea.SelectionChanged` and `TextArea.Changed` messages https://github.com/Textualize/textual/pull/3442
+- Added `wait_for_dismiss` parameter to `App.push_screen` https://github.com/Textualize/textual/pull/3477
 
 ### Changed
 

From 61ba093d1f0574ba3500473f04fc962b51ff2788 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 19:46:53 +0100
Subject: [PATCH 07/11] simplify

---
 src/textual/app.py | 29 +++++------------------------
 1 file changed, 5 insertions(+), 24 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index 288b1beaed..7e6c174722 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -247,25 +247,6 @@ def fileno(self) -> int:
         return -1
 
 
-class ScreenAwaitable(Generic[ScreenResultType]):
-    """An optional awaitable to get the result of a screen."""
-
-    def __init__(
-        self, await_mount: AwaitMount, future: asyncio.Future[ScreenResultType]
-    ) -> None:
-        self._await_mount = await_mount
-        self._future = future
-
-    def __await__(self) -> Generator[None, None, ScreenResultType]:
-        async def await_screen() -> ScreenResultType:
-            """Await the mount and the future."""
-            await self._await_mount
-            await self._future
-            return self._future.result()
-
-        return await_screen().__await__()
-
-
 @rich.repr.auto
 class App(Generic[ReturnType], DOMNode):
     """The base class for Textual Applications."""
@@ -1810,7 +1791,7 @@ def push_screen(
         screen: Screen[ScreenResultType] | str,
         callback: ScreenResultCallbackType[ScreenResultType] | None = None,
         wait_for_dismiss: Literal[True] = True,
-    ) -> ScreenAwaitable[ScreenResultType]:
+    ) -> asyncio.Future[ScreenResultType]:
         ...
 
     def push_screen(
@@ -1818,7 +1799,7 @@ def push_screen(
         screen: Screen[ScreenResultType] | str,
         callback: ScreenResultCallbackType[ScreenResultType] | None = None,
         wait_for_dismiss: bool = False,
-    ) -> AwaitMount | ScreenAwaitable[ScreenResultType]:
+    ) -> AwaitMount | asyncio.Future[ScreenResultType]:
         """Push a new [screen](/guide/screens) on the screen stack, making it the current screen.
 
         Args:
@@ -1831,7 +1812,8 @@ def push_screen(
             NoActiveWorker: If using `wait_for_dismiss` outside of a worker.
 
         Returns:
-            An optional awaitable that awaits the mounting of the screen and its children.
+            An optional awaitable that awaits the mounting of the screen and its children, or an asyncio Future
+                to await the result of the screen.
         """
         if not isinstance(screen, (Screen, str)):
             raise TypeError(
@@ -1868,8 +1850,7 @@ def push_screen(
                 raise NoActiveWorker(
                     "push_screen must be run from a worker when `wait_for_dismiss` is True"
                 ) from None
-            screen_awaitable = ScreenAwaitable(await_mount, future)
-            return screen_awaitable
+            return future
         else:
             return await_mount
 

From 77c8e7cbfba568b94f097e0c2908b46c1dcab888 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 19:49:55 +0100
Subject: [PATCH 08/11] typing

---
 src/textual/app.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index 7e6c174722..4a14a12e20 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -1780,8 +1780,8 @@ def _replace_screen(self, screen: Screen) -> Screen:
     def push_screen(
         self,
         screen: Screen[ScreenResultType] | str,
-        callback: ScreenResultCallbackType[ScreenResultType] | None = None,
-        wait_for_dismiss: Literal[False] = False,
+        callback: ScreenResultCallbackType[ScreenResultType] | None,
+        wait_for_dismiss: Literal[False],
     ) -> AwaitMount:
         ...
 
@@ -1789,8 +1789,8 @@ def push_screen(
     def push_screen(
         self,
         screen: Screen[ScreenResultType] | str,
-        callback: ScreenResultCallbackType[ScreenResultType] | None = None,
-        wait_for_dismiss: Literal[True] = True,
+        callback: ScreenResultCallbackType[ScreenResultType] | None,
+        wait_for_dismiss: Literal[True],
     ) -> asyncio.Future[ScreenResultType]:
         ...
 

From 6399dc9b0b3f498cf5d5746ef4aa490045f8ebe3 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 19:53:29 +0100
Subject: [PATCH 09/11] dot

---
 src/textual/screen.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/textual/screen.py b/src/textual/screen.py
index 89f26f504e..be09b66e7c 100644
--- a/src/textual/screen.py
+++ b/src/textual/screen.py
@@ -82,7 +82,7 @@ def __init__(
         Args:
             requester: The object making a request for the callback.
             callback: The callback function.
-            future: A Future to hold the result
+            future: A Future to hold the result.
         """
         self.requester = requester
         """The object in the DOM that requested the callback."""

From 1130fdac90dcf6dfb0c7535fa01adcd59e3c87d9 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sun, 8 Oct 2023 20:03:59 +0100
Subject: [PATCH 10/11] typing

---
 src/textual/app.py    | 8 ++++----
 tests/test_screens.py | 7 ++++---
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index 4a14a12e20..7e6c174722 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -1780,8 +1780,8 @@ def _replace_screen(self, screen: Screen) -> Screen:
     def push_screen(
         self,
         screen: Screen[ScreenResultType] | str,
-        callback: ScreenResultCallbackType[ScreenResultType] | None,
-        wait_for_dismiss: Literal[False],
+        callback: ScreenResultCallbackType[ScreenResultType] | None = None,
+        wait_for_dismiss: Literal[False] = False,
     ) -> AwaitMount:
         ...
 
@@ -1789,8 +1789,8 @@ def push_screen(
     def push_screen(
         self,
         screen: Screen[ScreenResultType] | str,
-        callback: ScreenResultCallbackType[ScreenResultType] | None,
-        wait_for_dismiss: Literal[True],
+        callback: ScreenResultCallbackType[ScreenResultType] | None = None,
+        wait_for_dismiss: Literal[True] = True,
     ) -> asyncio.Future[ScreenResultType]:
         ...
 
diff --git a/tests/test_screens.py b/tests/test_screens.py
index aa2a3c80a1..d715cd0a28 100644
--- a/tests/test_screens.py
+++ b/tests/test_screens.py
@@ -418,7 +418,7 @@ def on_mount(self):
 async def test_push_screen_wait_for_dismiss() -> None:
     """Test push_screen returns result."""
 
-    class QuitScreen(Screen):
+    class QuitScreen(Screen[bool]):
         BINDINGS = [
             ("y", "quit(True)"),
             ("n", "quit(False)"),
@@ -445,7 +445,7 @@ async def action_exit(self) -> None:
 
     results.clear()
     app = ScreensApp()
-    # Press X to exit, then Y to dismiss, expect True result
+    # Press X to exit, then Y to dismiss, expect False result
     async with app.run_test() as pilot:
         await pilot.press("x", "n")
     assert results == [False]
@@ -454,7 +454,7 @@ async def action_exit(self) -> None:
 async def test_push_screen_wait_for_dismiss_no_worker() -> None:
     """Test wait_for_dismiss raises NoActiveWorker when not using workers."""
 
-    class QuitScreen(Screen):
+    class QuitScreen(Screen[bool]):
         BINDINGS = [
             ("y", "quit(True)"),
             ("n", "quit(False)"),
@@ -473,6 +473,7 @@ async def action_exit(self) -> None:
             results.append(result)
 
     app = ScreensApp()
+    # using `wait_for_dismiss` outside of a worker should raise NoActiveWorker
     with pytest.raises(NoActiveWorker):
         async with app.run_test() as pilot:
             await pilot.press("x", "y")

From 5133d5f4d08138bfddbe8f43a9843092a4ed5f9b Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Mon, 9 Oct 2023 11:37:16 +0100
Subject: [PATCH 11/11] Update tests/test_screens.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Rodrigo Girão Serrão <5621605+rodrigogiraoserrao@users.noreply.github.com>
---
 tests/test_screens.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_screens.py b/tests/test_screens.py
index d715cd0a28..e5ddacb4d7 100644
--- a/tests/test_screens.py
+++ b/tests/test_screens.py
@@ -445,7 +445,7 @@ async def action_exit(self) -> None:
 
     results.clear()
     app = ScreensApp()
-    # Press X to exit, then Y to dismiss, expect False result
+    # Press X to exit, then N to dismiss, expect False result
     async with app.run_test() as pilot:
         await pilot.press("x", "n")
     assert results == [False]