Preserve original sender when using call_after_refresh (#2806)

* Defer sender in call_after_refresh

* docstring

* test

* typing

Author: willmcgugan

CHANGELOG.md

@@ -22,6 +22,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 - Tooltips are now inherited, so will work with compound widgets
 
+### Fixed
+
+- call_after_refresh will preserve the sender within the callback https://github.com/Textualize/textual/pull/2806
+
 ## [0.28.0] - 2023-06-19
 
 ### Added

src/textual/message.py

@@ -49,7 +49,7 @@ def __init__(self) -> None:
 
     def __post_init__(self) -> None:
         """Allow dataclasses to initialize the object."""
-        self._sender: MessageTarget | None = active_message_pump.get(None)
+        self._sender: MessagePump | None = active_message_pump.get(None)
         self.time: float = _time.get_time()
         self._forwarded = False
         self._no_default_action = False
@@ -91,7 +91,7 @@ def _set_forwarded(self) -> None:
         """Mark this event as being forwarded."""
         self._forwarded = True
 
-    def _set_sender(self, sender: MessageTarget) -> None:
+    def _set_sender(self, sender: MessagePump) -> None:
         """Set the sender."""
         self._sender = sender
 

src/textual/message_pump.py

@@ -422,7 +422,9 @@ def call_next(self, callback: Callback, *args: Any, **kwargs: Any) -> None:
 
     def _on_invoke_later(self, message: messages.InvokeLater) -> None:
         # Forward InvokeLater message to the Screen
-        self.app.screen._invoke_later(message.callback)
+        self.app.screen._invoke_later(
+            message.callback, message._sender or active_message_pump.get()
+        )
 
     def _close_messages_no_wait(self) -> None:
         """Request the message queue to immediately exit."""

src/textual/screen.py

@@ -27,7 +27,7 @@
 from . import errors, events, messages
 from ._callback import invoke
 from ._compositor import Compositor, MapGeometry
-from ._context import visible_screen_stack
+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 .binding import Binding
@@ -48,6 +48,7 @@
 
     # Unused & ignored imports are needed for the docs to link to these objects:
     from .errors import NoWidget  # type: ignore  # noqa: F401
+    from .message_pump import MessagePump
 
 # Screen updates will be batched so that they don't happen more often than 60 times per second:
 UPDATE_PERIOD: Final[float] = 1 / 60
@@ -154,7 +155,7 @@ def __init__(
         self._compositor = Compositor()
         self._dirty_widgets: set[Widget] = set()
         self.__update_timer: Timer | None = None
-        self._callbacks: list[CallbackType] = []
+        self._callbacks: list[tuple[CallbackType, MessagePump]] = []
         self._result_callbacks: list[ResultCallback[ScreenResultType]] = []
 
         self._tooltip_widget: Widget | None = None
@@ -587,17 +588,22 @@ async def _invoke_and_clear_callbacks(self) -> None:
         if self._callbacks:
             callbacks = self._callbacks[:]
             self._callbacks.clear()
-            for callback in callbacks:
-                await invoke(callback)
+            for callback, message_pump in callbacks:
+                reset_token = active_message_pump.set(message_pump)
+                try:
+                    await invoke(callback)
+                finally:
+                    active_message_pump.reset(reset_token)
 
-    def _invoke_later(self, callback: CallbackType) -> None:
+    def _invoke_later(self, callback: CallbackType, sender: MessagePump) -> None:
         """Enqueue a callback to be invoked after the screen is repainted.
 
         Args:
             callback: A callback.
+            sender: The sender (active message pump) of the callback.
         """
 
-        self._callbacks.append(callback)
+        self._callbacks.append((callback, sender))
         self.check_idle()
 
     def _push_result_callback(

tests/test_call_later.py

@@ -1,6 +1,8 @@
 import asyncio
 
+from textual._context import active_message_pump
 from textual.app import App
+from textual.message_pump import MessagePump
 
 
 class CallLaterApp(App[None]):
@@ -30,13 +32,19 @@ async def test_call_after_refresh() -> None:
 
     called_event = asyncio.Event()
 
+    callback_message_pump: MessagePump | None = None
+
     def callback() -> None:
         nonlocal display_count
+        nonlocal callback_message_pump
         called_event.set()
         display_count = app.display_count
+        callback_message_pump = active_message_pump.get()
 
     async with app.run_test():
         app.call_after_refresh(callback)
         await asyncio.wait_for(called_event.wait(), 1)
         app_display_count = app.display_count
         assert app_display_count == display_count
+        # Should be app because the callback should be in the same context as app
+        assert callback_message_pump is app