From ee71b34459a9d9505236fd0f35bd5f604f25c159 Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Mon, 24 Jun 2024 14:28:18 +0100 Subject: [PATCH 1/4] await screens --- src/textual/_worker_manager.py | 6 ++- src/textual/app.py | 73 ++++++++++++++++---------- src/textual/await_complete.py | 12 +++++ src/textual/message_pump.py | 5 +- src/textual/screen.py | 22 ++++---- src/textual/widget.py | 2 +- src/textual/widgets/_directory_tree.py | 2 +- 7 files changed, 80 insertions(+), 42 deletions(-) diff --git a/src/textual/_worker_manager.py b/src/textual/_worker_manager.py index ab69947718..6c90b9641a 100644 --- a/src/textual/_worker_manager.py +++ b/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 diff --git a/src/textual/app.py b/src/textual/app.py index f34ca20370..d5a41943b7 100644 --- a/src/textual/app.py +++ b/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,16 @@ 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(): + 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 +2037,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 +2053,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 +2158,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: @@ -2164,19 +2172,23 @@ def switch_screen(self, screen: Screen | str) -> AwaitMount: f"switch_screen requires a Screen instance or str; not {screen!r}" ) - 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, []) + async def do_switch() -> None: + 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 - previous_screen = self._replace_screen(self._screen_stack.pop()) - previous_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 + await await_mount() + + previous_screen = await self._replace_screen(self._screen_stack.pop()) + previous_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 AwaitComplete(do_switch()).call_next(self) def install_screen(self, screen: Screen, name: str) -> None: """Install a screen. @@ -2238,22 +2250,26 @@ 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._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: + previous_screen = await self._replace_screen(screen_stack.pop()) + previous_screen._pop_result_callback() + self.screen.post_message(events.ScreenResume()) + self.log.system(f"{self.screen} is active") + + 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. @@ -2354,6 +2370,7 @@ def _handle_exception(self, error: Exception) -> None: Args: error: An exception instance. """ + self.log.error(error) self._return_code = 1 # If we're running via pilot and this is the first exception encountered, # take note of it so that we can re-raise for test frameworks later. @@ -3408,7 +3425,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. diff --git a/src/textual/await_complete.py b/src/textual/await_complete.py index 2ff1068862..ebc3c84ea5 100644 --- a/src/textual/await_complete.py +++ b/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 diff --git a/src/textual/message_pump.py b/src/textual/message_pump.py index a91255a3f4..7ee67fd85e 100644 --- a/src/textual/message_pump.py +++ b/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.""" diff --git a/src/textual/screen.py b/src/textual/screen.py index b7c6b9123e..5ca9070855 100644 --- a/src/textual/screen.py +++ b/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,13 +1227,15 @@ 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`. @@ -1240,22 +1243,22 @@ def dismiss(self, result: ScreenResultType | Type[_NoResult] = _NoResult) -> boo 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. @@ -1263,6 +1266,7 @@ def action_dismiss( 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: diff --git a/src/textual/widget.py b/src/textual/widget.py index 5d14903d3d..feeda84e93 100644 --- a/src/textual/widget.py +++ b/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 diff --git a/src/textual/widgets/_directory_tree.py b/src/textual/widgets/_directory_tree.py index 16f9d0a100..fdc1efcee7 100644 --- a/src/textual/widgets/_directory_tree.py +++ b/src/textual/widgets/_directory_tree.py @@ -170,7 +170,7 @@ def _add_to_load_queue(self, node: TreeNode[DirEntry]) -> AwaitComplete: node.data.loaded = True self._load_queue.put_nowait(node) - return AwaitComplete(self._load_queue.join()) + return AwaitComplete(self, self._load_queue.join()) def reload(self) -> AwaitComplete: """Reload the `DirectoryTree` contents. From b40387d667898582eff99fc6f0ed630534d1267a Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Mon, 24 Jun 2024 16:01:05 +0100 Subject: [PATCH 2/4] polish --- CHANGELOG.md | 8 ++++++++ src/textual/app.py | 6 ++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 05e08c0fd0..ebc4f1eca5 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/textual/app.py b/src/textual/app.py index d5a41943b7..01b66ab8a8 100644 --- a/src/textual/app.py +++ b/src/textual/app.py @@ -1938,7 +1938,8 @@ def remove_mode(self, mode: str) -> AwaitComplete: stack = self._screen_stacks[mode] del self._screen_stacks[mode] - async def remove_screens(): + async def remove_screens() -> None: + """Remove screens.""" for screen in reversed(stack): await self._replace_screen(screen) @@ -2173,6 +2174,7 @@ def switch_screen(self, screen: Screen | str) -> AwaitComplete: ) async def do_switch() -> None: + """Task to perform switch.""" 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.") @@ -2264,6 +2266,7 @@ def pop_screen(self) -> AwaitComplete: ) async def do_pop() -> None: + """Task to pop the screen.""" previous_screen = await self._replace_screen(screen_stack.pop()) previous_screen._pop_result_callback() self.screen.post_message(events.ScreenResume()) @@ -2370,7 +2373,6 @@ def _handle_exception(self, error: Exception) -> None: Args: error: An exception instance. """ - self.log.error(error) self._return_code = 1 # If we're running via pilot and this is the first exception encountered, # take note of it so that we can re-raise for test frameworks later. From e93b3facb0969b765abaa741ad0c7360611af51f Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Mon, 24 Jun 2024 16:01:31 +0100 Subject: [PATCH 3/4] fix --- src/textual/widgets/_directory_tree.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/textual/widgets/_directory_tree.py b/src/textual/widgets/_directory_tree.py index fdc1efcee7..16f9d0a100 100644 --- a/src/textual/widgets/_directory_tree.py +++ b/src/textual/widgets/_directory_tree.py @@ -170,7 +170,7 @@ def _add_to_load_queue(self, node: TreeNode[DirEntry]) -> AwaitComplete: node.data.loaded = True self._load_queue.put_nowait(node) - return AwaitComplete(self, self._load_queue.join()) + return AwaitComplete(self._load_queue.join()) def reload(self) -> AwaitComplete: """Reload the `DirectoryTree` contents. From fe82774a465b1535e8cafad44b71b0db7a2f7da4 Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Mon, 24 Jun 2024 17:09:14 +0100 Subject: [PATCH 4/4] test fixes --- src/textual/app.py | 55 +++++++++++++++++++++--------------- tests/css/test_screen_css.py | 2 +- tests/test_screens.py | 16 ++++++----- 3 files changed, 43 insertions(+), 30 deletions(-) diff --git a/src/textual/app.py b/src/textual/app.py index 01b66ab8a8..4b667eb643 100644 --- a/src/textual/app.py +++ b/src/textual/app.py @@ -2173,22 +2173,25 @@ def switch_screen(self, screen: Screen | str) -> AwaitComplete: f"switch_screen requires a Screen instance or str; not {screen!r}" ) + 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 AwaitComplete.nothing() + + 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)") + async def do_switch() -> None: """Task to perform switch.""" - 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 await await_mount() - - previous_screen = await self._replace_screen(self._screen_stack.pop()) - previous_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)") + await self._replace_screen(top_screen) return AwaitComplete(do_switch()).call_next(self) @@ -2265,12 +2268,14 @@ def pop_screen(self) -> AwaitComplete: "Can't pop screen; there must be at least one screen on the stack" ) + previous_screen = screen_stack.pop() + previous_screen._pop_result_callback() + self.screen.post_message(events.ScreenResume()) + self.log.system(f"{self.screen} is active") + async def do_pop() -> None: """Task to pop the screen.""" - previous_screen = await self._replace_screen(screen_stack.pop()) - previous_screen._pop_result_callback() - self.screen.post_message(events.ScreenResume()) - self.log.system(f"{self.screen} is active") + await self._replace_screen(previous_screen) return AwaitComplete(do_pop()).call_next(self) @@ -3326,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. diff --git a/tests/css/test_screen_css.py b/tests/css/test_screen_css.py index 54138fb8a5..61d26a6a38 100644 --- a/tests/css/test_screen_css.py +++ b/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): diff --git a/tests/test_screens.py b/tests/test_screens.py index f2d51331e9..aef83dfdb4 100644 --- a/tests/test_screens.py +++ b/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():