Swap the result callbacks on screens to be a stack of callbacks

It is possible for the same instance of a screen to get pushed onto the
screen stack multiple times; as such we really need to keep track of all the
callback requests.

So here I register a callback for every screen push and clean it up on every
screen pop; with those without callbacks being no-ops.

src/textual/app.py

@@ -92,7 +92,7 @@ from .keys import (
 from .messages import CallbackType
 from .reactive import Reactive
 from .renderables.blank import Blank
-from .screen import Screen, ScreenResultCallbackType
+from .screen import Screen, ScreenResultCallbackType, ScreenResultType
 from .widget import AwaitMount, Widget
 
 if TYPE_CHECKING:
@@ -1380,13 +1380,18 @@ class App(Generic[ReturnType], DOMNode):
         return screen
 
     def push_screen(
-        self, screen: Screen | str, callback: ScreenResultCallbackType | None = None
+        self,
+        screen: Screen[ScreenResultType] | str,
+        callback: ScreenResultCallbackType[ScreenResultType] | None = None,
     ) -> AwaitMount:
         """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 is called if the screen is dismissed with a result.
 
+        Returns:
+            An awaitable that awaits the mounting of the screen and its children.
         """
         if not isinstance(screen, (Screen, str)):
             raise TypeError(
@@ -1397,7 +1402,7 @@ class App(Generic[ReturnType], DOMNode):
             self.screen.post_message(events.ScreenSuspend())
             self.screen.refresh()
         next_screen, await_mount = self._get_screen(screen)
-        next_screen._result_callback = callback
+        next_screen._push_result_callback(self.screen, callback)
         self._screen_stack.append(next_screen)
         next_screen.post_message(events.ScreenResume())
         self.log.system(f"{self.screen} is current (PUSHED)")
@@ -1496,6 +1501,7 @@ class App(Generic[ReturnType], DOMNode):
                 "Can't pop screen; there must be at least one screen on the stack"
             )
         previous_screen = self._replace_screen(screen_stack.pop())
+        previous_screen._pop_result_callback()
         self.screen._screen_resized(self.size)
         self.screen.post_message(events.ScreenResume())
         self.log.system(f"{self.screen} is active")

src/textual/screen.py

@@ -6,7 +6,6 @@ The `Screen` class is a special widget which represents the content in the termi
 
 from __future__ import annotations
 
-from inspect import isawaitable
 from typing import (
     TYPE_CHECKING,
     Awaitable,
@@ -53,6 +52,38 @@ ScreenResultCallbackType = Union[
 """Type of a screen result callback function."""
 
 
+class ResultCallback(Generic[ScreenResultType]):
+    """Holds the details of a callback."""
+
+    def __init__(
+        self,
+        requester: Widget,
+        callback: ScreenResultCallbackType[ScreenResultType] | None,
+    ) -> None:
+        """Initialise the result callback object.
+
+        Args:
+            requester: The object making a request for the callback.
+            callback: The callback function.
+        """
+        self.requester: Widget = requester
+        """The object in the DOM that requested the callback."""
+        self.callback: ScreenResultCallbackType | None = callback
+        """The callback function."""
+
+    def __call__(self, result: ScreenResultType) -> None:
+        """Call the callback, passing the given result.
+
+        Args:
+            result: The result to pass to the callback.
+
+        Note:
+            If the callback is `None` this will be a no-op.
+        """
+        if self.callback is not None:
+            self.requester.call_next(self.callback, result)
+
+
 @rich.repr.auto
 class Screen(Generic[ScreenResultType], Widget):
     """The base class for screens."""
@@ -95,7 +126,7 @@ class Screen(Generic[ScreenResultType], Widget):
         self._dirty_widgets: set[Widget] = set()
         self.__update_timer: Timer | None = None
         self._callbacks: list[CallbackType] = []
-        self._result_callback: ScreenResultCallbackType | None = None
+        self._result_callbacks: list[ResultCallback[ScreenResultType]] = []
 
     @property
     def is_modal(self) -> bool:
@@ -468,6 +499,25 @@ class Screen(Generic[ScreenResultType], Widget):
         self._callbacks.append(callback)
         self.check_idle()
 
+    def _push_result_callback(
+        self,
+        requester: Widget,
+        callback: ScreenResultCallbackType[ScreenResultType] | None,
+    ) -> None:
+        """Add a result callback to the screen.
+
+        Args:
+            requester: The object requesting the callback.
+            callback: The callback.
+        """
+        self._result_callbacks.append(
+            ResultCallback[ScreenResultType](requester, callback)
+        )
+
+    def _pop_result_callback(self) -> None:
+        """Remove the latest result callback from the stack."""
+        self._result_callbacks.pop()
+
     def _refresh_layout(
         self, size: Size | None = None, full: bool = False, scroll: bool = False
     ) -> None:
@@ -663,10 +713,24 @@ class Screen(Generic[ScreenResultType], Widget):
             self.post_message(event)
 
     def dismiss(self) -> None:
-        """Dismiss the screen."""
+        """Dismiss the screen.
+
+        This is a convenience wrapper around
+        [`App.pop_screen`][textual.app.App.pop_screen].
+
+        Note:
+            If the screen was pushed with a callback it will *not* be called
+            when using `dismiss`. Callbacks are only called when using
+            [`dismiss_with`][textual.screen.Screen.dismiss_with].
+        """
         self.app.pop_screen()
 
     def action_dismiss(self) -> None:
+        """A wrapper around [`dismiss`][textual.screen.Screen.dismiss] that can be called as an action.
+
+        This allows for easily dismissing a screen from a keyboard binding,
+        for example.
+        """
         self.dismiss()
 
     def dismiss_with(self, result: ScreenResultType) -> None:
@@ -674,16 +738,31 @@ class Screen(Generic[ScreenResultType], Widget):
 
         Args:
             result: The result to be passed to the result callback.
+
+        Note:
+            If the screen was pushed with a callback, the callback will be
+            called with the given result and then a call to
+            [`App.pop_screen`][textual.app.App.pop_screen] is performed. If
+            no callback was provided calling this method is the same as
+            simply calling [`dismiss`][textual.screen.Screen.dismiss].
         """
-        if self._result_callback is not None:
-            try:
-                callback_screen = self.app.screen_stack[-2]
-            except IndexError:
-                callback_screen = self
-            callback_screen.call_next(self._result_callback, result)
+        if self._result_callbacks:
+            self._result_callbacks[-1](result)
         self.dismiss()
 
     def action_dismiss_with(self, result: ScreenResultType) -> None:
+        """A wrapper around [`dismiss_with`][textual.screen.Screen.dismiss_with] that can be called as an action.
+
+        Args:
+            result: The result to be passed to the result callback.
+
+        Note:
+            If the screen was pushed with a callback, the callback will be
+            called with the given result and then a call to
+            [`App.pop_screen`][textual.app.App.pop_screen] is performed. If
+            no callback was provided calling this method is the same as
+            simply calling [`dismiss`][textual.screen.Screen.dismiss].
+        """
         self.dismiss_with(result)