Try and better settle focus after a focused widget is removed

sandbox/davep/single_widget.py

@@ -0,0 +1,35 @@
+"""Sandbox app to test widget removal and focus-handing.
+
+Written to help test changes made to solve:
+
+  https://github.com/Textualize/textual/issues/939
+"""
+
+from textual.app import App
+from textual.containers import Container
+from textual.widgets import Static
+
+
+class NonFocusParent(Static):
+    def compose(self):
+        yield Static("Test")
+
+
+class SingleWidgetTester(App[None]):
+
+    BINDINGS = [("a", "add_widget", "Add Widget"), ("d", "del_widget", "Delete Widget")]
+
+    def compose(self):
+        yield Container()
+
+    def action_add_widget(self):
+        self.query_one(Container).mount(NonFocusParent())
+
+    def action_del_widget(self):
+        candidates = self.query(NonFocusParent)
+        if candidates:
+            candidates.last().remove()
+
+
+if __name__ == "__main__":
+    SingleWidgetTester().run()

src/textual/app.py

@@ -1419,14 +1419,37 @@ async def _on_remove(self, event: events.Remove) -> None:
         widget = event.widget
         parent = widget.parent
 
-        widget.reset_focus()
-
-        remove_widgets = widget.walk_children(
-            Widget, with_self=True, method="depth", reverse=True
+        # Get the list of all the widgets that are about to be removed.
+        remove_widgets = list(
+            widget.walk_children(Widget, with_self=True, method="depth", reverse=True)
         )
+
+        # Of the list of widgets to be removed, let's get a set of those
+        # that can receive focus -- we're about to go on a focus hunt and
+        # this can be a subset of the above (that is: it's possible that not
+        # every widget we're about to remove is one that can hold focus).
+        remove_focusable = {widget for widget in remove_widgets if widget.can_focus}
+
+        # Assume that we aren't going to move focus at all.
+        settle_focus = self.screen.focused
+
+        # If there are more items in the focus chain than there are items to
+        # be removed that can receive focus...
+        if len(widget.screen.focus_chain) > len(remove_focusable):
+            # ...go and seek the previous widget in the focus chain that
+            # isn't going to be removed.
+            while settle_focus is not None and settle_focus in remove_widgets:
+                settle_focus = self.screen.focus_from(settle_focus, -1)
+        else:
+            # It seems we can't settle on anywhere to focus.
+            settle_focus = None
+
         for child in remove_widgets:
             await child._close_messages()
             self._unregister(child)
+
+        self.screen.set_focus(settle_focus)
+
         if parent is not None:
             parent.refresh(layout=True)
 

src/textual/screen.py

@@ -175,37 +175,52 @@ def focus_chain(self) -> list[Widget]:
 
         return widgets
 
-    def _move_focus(self, direction: int = 0) -> Widget | None:
-        """Move the focus in the given direction.
+    def focus_from(self, widget: Widget | None, direction: int = 0) -> Widget | None:
+        """Locate a candidate widget given a focus direction.
 
         Args:
-            direction (int, optional): 1 to move forward, -1 to move backward, or
+            direction (int, optional): 1 to locate forward, -1 to locate backward, or
                 0 to keep the current focus.
 
         Returns:
-            Widget | None: Newly focused widget, or None for no focus.
+            Widget | None: The candidate widget, or None for no focus.
         """
         focusable_widgets = self.focus_chain
 
         if not focusable_widgets:
             # Nothing focusable, so nothing to do
-            return self.focused
-        if self.focused is None:
-            # Nothing currently focused, so focus the first one
-            self.set_focus(focusable_widgets[0])
+            return None
+        if widget is None:
+            # No widget given so the first in the chain is the best
+            # candidate from here.
+            return focusable_widgets[0]
         else:
             try:
-                # Find the index of the currently focused widget
-                current_index = focusable_widgets.index(self.focused)
+                # Find the index of the widget to seek from.
+                current_index = focusable_widgets.index(widget)
             except ValueError:
-                # Focused widget was removed in the interim, start again
-                self.set_focus(focusable_widgets[0])
+                # That widget was removed in the interim, start again
+                return focusable_widgets[0]
             else:
-                # Only move the focus if we are currently showing the focus
+                # Only seek a new widget if we've been given a direction.
                 if direction:
-                    current_index = (current_index + direction) % len(focusable_widgets)
-                    self.set_focus(focusable_widgets[current_index])
+                    return focusable_widgets[
+                        (current_index + direction) % len(focusable_widgets)
+                    ]
+
+        return None
+
+    def _move_focus(self, direction: int = 0) -> Widget | None:
+        """Move the focus in the given direction.
+
+        Args:
+            direction (int, optional): 1 to move forward, -1 to move backward, or
+                0 to keep the current focus.
 
+        Returns:
+            Widget | None: Newly focused widget, or None for no focus.
+        """
+        self.set_focus(self.focus_from(self.focused, direction))
         return self.focused
 
     def focus_next(self) -> Widget | None:

src/textual/widget.py

@@ -1775,7 +1775,6 @@ def refresh(
 
     def remove(self) -> None:
         """Remove the Widget from the DOM (effectively deleting it)"""
-        self.display = False
         self.app.post_message_no_wait(events.Remove(self, widget=self))
 
     def render(self) -> RenderableType: