prune

CHANGELOG.md

@@ -5,6 +5,12 @@ 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
+
+- More predictable DOM removals. https://github.com/Textualize/textual/pull/4708
+
 ## [0.71.0] - 2024-06-29
 
 ### Changed

src/textual/app.py

@@ -99,11 +99,10 @@
     _get_key_display,
     _get_unicode_name_from_key,
 )
-from .messages import CallbackType
+from .messages import CallbackType, Prune
 from .notifications import Notification, Notifications, Notify, SeverityLevel
 from .reactive import Reactive
 from .renderables.blank import Blank
-from .rlock import RLock
 from .screen import (
     ActiveBinding,
     Screen,
@@ -582,7 +581,6 @@ def __init__(
             else None
         )
         self._screenshot: str | None = None
-        self._dom_lock = RLock()
         self._dom_ready = False
         self._batch_count = 0
         self._notifications = Notifications()
@@ -1458,6 +1456,11 @@ def on_app_ready() -> None:
             app_ready_event.set()
 
         async def run_app(app: App) -> None:
+            """Run the apps message loop.
+
+            Args:
+                app: App to run.
+            """
             if message_hook is not None:
                 message_hook_context_var.set(message_hook)
             app._loop = asyncio.get_running_loop()
@@ -2743,6 +2746,9 @@ def _register(
 
         apply_stylesheet = self.stylesheet.apply
         for widget in widget_list:
+            widget._closing = False
+            widget._closed = False
+            widget._pruning = False
             if not isinstance(widget, Widget):
                 raise AppError(f"Can't register {widget!r}; expected a Widget instance")
             if widget not in self._registry:
@@ -2799,25 +2805,23 @@ def is_mounted(self, widget: Widget) -> bool:
     async def _close_all(self) -> None:
         """Close all message pumps."""
 
-        async with self._dom_lock:
-            # Close all screens on all stacks:
-            for stack in self._screen_stacks.values():
-                for stack_screen in reversed(stack):
-                    if stack_screen._running:
-                        await self._prune_node(stack_screen)
-                stack.clear()
-
-            # Close pre-defined screens.
-            for screen in self.SCREENS.values():
-                if isinstance(screen, Screen) and screen._running:
-                    await self._prune_node(screen)
+        # Close all screens on all stacks:
+        for stack in self._screen_stacks.values():
+            for stack_screen in reversed(stack):
+                if stack_screen._running:
+                    await self._prune(stack_screen)
+            stack.clear()
 
-            # Close any remaining nodes
-            # Should be empty by now
-            remaining_nodes = list(self._registry)
+        # Close pre-defined screens.
+        for screen in self.SCREENS.values():
+            if isinstance(screen, Screen) and screen._running:
+                await self._prune(screen)
 
-            for child in remaining_nodes:
-                await child._close_messages()
+        # Close any remaining nodes
+        # Should be empty by now
+        remaining_nodes = list(self._registry)
+        for child in remaining_nodes:
+            await child._close_messages()
 
     async def _shutdown(self) -> None:
         self._begin_batch()  # Prevents any layout / repaint while shutting down
@@ -3300,189 +3304,53 @@ async def _on_app_blur(self, event: events.AppBlur) -> None:
         self.app_focus = False
         self.screen.refresh_bindings()
 
-    def _detach_from_dom(self, widgets: list[Widget]) -> list[Widget]:
-        """Detach a list of widgets from the DOM.
+    def _prune(self, *nodes: Widget, parent: DOMNode | None = None) -> AwaitRemove:
+        """Prune nodes from DOM.
 
         Args:
-            widgets: The list of widgets to detach from the DOM.
+            parent: Parent node.
 
         Returns:
-            The list of widgets that should be pruned.
-
-        Note:
-            A side-effect of calling this function is that each parent of
-            each affected widget will be made to forget about the affected
-            child.
-        """
-
-        # We've been given a list of widgets to remove, but removing those
-        # will also result in other (descendent) widgets being removed. So
-        # to start with let's get a list of everything that's not going to
-        # be in the DOM by the time we've finished. Note that, at this
-        # point, it's entirely possible that there will be duplicates.
-        everything_to_remove: list[Widget] = []
-        for widget in widgets:
-            everything_to_remove.extend(
-                widget.walk_children(
-                    Widget, with_self=True, method="depth", reverse=True
-                )
-            )
+            Optional awaitable.
+        """
+        if not nodes:
+            return AwaitRemove([])
+        pruning_nodes: set[Widget] = {*nodes}
+        for node in nodes:
+            node.post_message(Prune())
+            pruning_nodes.update(node.walk_children(with_self=True))
 
-        # Next up, let's quickly create a deduped collection of things to
-        # 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)
         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:
+            screen = nodes[0].screen
+        except (ScreenStackError, NoScreen):
             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.
-        # In other words: find the smallest set of ancestors in the DOM that
-        # will remove the widgets requested for removal, and also ensure
-        # that all knock-on effects happen too.
-        request_remove = set(widgets)
-        pruned_remove = [
-            widget for widget in widgets if request_remove.isdisjoint(widget.ancestors)
-        ]
-
-        # Now that we know that minimal set of widgets, we go through them
-        # and get their parents to forget about them. This has the effect of
-        # snipping each affected branch from the DOM.
-        for widget in pruned_remove:
-            if widget.parent is not None:
-                widget.parent._nodes._remove(widget)
-
-        for node in pruned_remove:
-            node._detach()
-
-        # Return the list of widgets that should end up being sent off in a
-        # prune event.
-        return pruned_remove
-
-    def _walk_children(self, root: Widget) -> Iterable[list[Widget]]:
-        """Walk children depth first, generating widgets and a list of their siblings.
-
-        Returns:
-            The child widgets of root.
-        """
-        stack: list[Widget] = [root]
-        pop = stack.pop
-        push = stack.append
-
-        while stack:
-            widget = pop()
-            children = [*widget._nodes, *widget._get_virtual_dom()]
-            if children:
-                yield children
-            for child in widget._nodes:
-                push(child)
-
-    def _remove_nodes(
-        self, widgets: list[Widget], parent: DOMNode | None
-    ) -> AwaitRemove:
-        """Remove nodes from DOM, and return an awaitable that awaits cleanup.
-
-        Args:
-            widgets: List of nodes to remove.
-            parent: Parent node of widgets, or None for no parent.
+        else:
+            if screen.focused and screen.focused in pruning_nodes:
+                screen._reset_focus(screen.focused, list(pruning_nodes))
 
-        Returns:
-            Awaitable that returns when the nodes have been fully removed.
-        """
+        for node in pruning_nodes:
+            node._pruning = True
 
-        async def prune_widgets_task(
-            widgets: list[Widget], finished_event: asyncio.Event
-        ) -> None:
-            """Prune widgets as a background task.
+        def post_mount() -> None:
+            """Called after removing children."""
 
-            Args:
-                widgets: Widgets to prune.
-                finished_event: Event to set when complete.
-            """
-            try:
-                await self._prune_nodes(widgets)
-            finally:
-                finished_event.set()
+            if parent is not None:
                 try:
-                    self._update_mouse_over(self.screen)
-                except ScreenStackError:
+                    screen = parent.screen
+                except (ScreenStackError, NoScreen):
                     pass
-                if parent is not None:
+                else:
+                    if screen._running:
+                        self._update_mouse_over(screen)
+                finally:
                     parent.refresh(layout=True)
 
-        removed_widgets = self._detach_from_dom(widgets)
-
-        finished_event = asyncio.Event()
-        remove_task = create_task(
-            prune_widgets_task(removed_widgets, finished_event), name="prune nodes"
+        await_complete = AwaitRemove(
+            [task for node in nodes if (task := node._task) is not None],
+            post_mount,
         )
-
-        await_remove = AwaitRemove(finished_event, remove_task)
-        self.call_next(await_remove)
-        return await_remove
-
-    async def _prune_nodes(self, widgets: list[Widget]) -> None:
-        """Remove nodes and children.
-
-        Args:
-            widgets: Widgets to remove.
-        """
-
-        for widget in widgets:
-            async with self._dom_lock:
-                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.
-
-        Args:
-            root: Node to remove.
-        """
-        # Pruning a node that has been removed is a no-op
-
-        if root not in self._registry:
-            return
-
-        node_children = list(self._walk_children(root))
-
-        for children in reversed(node_children):
-            # Closing children can be done asynchronously.
-            close_children = [
-                child for child in children if child._running and not child._closing
-            ]
-
-            # TODO: What if a message pump refuses to exit?
-            if close_children:
-                close_messages = [
-                    child._close_messages(wait=True) for child in close_children
-                ]
-                try:
-                    # Close all the children
-                    await asyncio.wait_for(
-                        asyncio.gather(*close_messages), self.CLOSE_TIMEOUT
-                    )
-                except asyncio.TimeoutError:
-                    # Likely a deadlock if we get here
-                    # If not a deadlock, increase CLOSE_TIMEOUT, or set it to None
-                    raise asyncio.TimeoutError(
-                        f"Timeout waiting for {close_children!r} to close; possible deadlock (consider changing App.CLOSE_TIMEOUT)\n"
-                    ) from None
-                finally:
-                    for child in children:
-                        self._unregister(child)
-
-        await root._close_messages(wait=True)
-        self._unregister(root)
+        self.call_next(await_complete)
+        return await_complete
 
     def _watch_app_focus(self, focus: bool) -> None:
         """Respond to changes in app focus."""

src/textual/await_remove.py

@@ -2,33 +2,36 @@
 An *optionally* awaitable object returned by methods that remove widgets.
 """
 
-from asyncio import Event, Task
-from typing import Generator
+from __future__ import annotations
 
+import asyncio
+from asyncio import Task, gather
+from typing import Generator
 
-class AwaitRemove:
-    """An awaitable returned by a method that removes DOM nodes.
+from ._callback import invoke
+from ._types import CallbackType
 
-    Returned by [Widget.remove][textual.widget.Widget.remove] and
-    [DOMQuery.remove][textual.css.query.DOMQuery.remove].
-    """
 
-    def __init__(self, finished_flag: Event, task: Task) -> None:
-        """Initialise the instance of ``AwaitRemove``.
+class AwaitRemove:
+    """An awaitable that waits for nodes to be removed."""
 
-        Args:
-            finished_flag: The asyncio event to wait on.
-            task: The task which does the remove (required to keep a reference).
-        """
-        self.finished_flag = finished_flag
-        self._task = task
+    def __init__(
+        self, tasks: list[Task], post_remove: CallbackType | None = None
+    ) -> None:
+        self._tasks = tasks
+        self._post_remove = post_remove
 
     async def __call__(self) -> None:
         await self
 
     def __await__(self) -> Generator[None, None, None]:
+        current_task = asyncio.current_task()
+        tasks = [task for task in self._tasks if task is not current_task]
+
         async def await_prune() -> None:
             """Wait for the prune operation to finish."""
-            await self.finished_flag.wait()
+            await gather(*tasks)
+            if self._post_remove is not None:
+                await invoke(self._post_remove)
 
         return await_prune().__await__()

src/textual/css/query.py

@@ -410,8 +410,7 @@ def remove(self) -> AwaitRemove:
             An awaitable object that waits for the widgets to be removed.
         """
         app = active_app.get()
-        await_remove = app._remove_nodes(list(self), self._node)
-        return await_remove
+        return app._prune(*self.nodes, parent=self._node)
 
     def set_styles(
         self, css: str | None = None, **update_styles