Try and better settle focus after a focused widget is removed (redux)

sandbox/davep/Makefile

@@ -0,0 +1,2 @@
+all:
+	poetry run textual run --dev focus_removal_tester.py

sandbox/davep/focus_removal_tester.py

@@ -0,0 +1,45 @@
+"""Focus removal tester.
+
+https://github.com/Textualize/textual/issues/939
+"""
+
+from textual.app import App
+from textual.containers import Container
+from textual.widgets import Static, Header, Footer, Button
+
+
+class LeftButton(Button):
+    pass
+
+
+class RightButton(Button):
+    pass
+
+
+class NonFocusParent(Static):
+    def compose(self):
+        yield LeftButton("Do Not Press")
+        yield Static("Test")
+        yield RightButton("Really Do Not Press")
+
+
+class FocusRemovalTester(App[None]):
+
+    BINDINGS = [("a", "add_widget", "Add Widget"), ("d", "del_widget", "Delete Widget")]
+
+    def compose(self):
+        yield Header()
+        yield Container()
+        yield Footer()
+
+    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__":
+    FocusRemovalTester().run()

src/textual/app.py

@@ -1419,11 +1419,16 @@ 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
+        remove_widgets = list(
+            widget.walk_children(Widget, with_self=True, method="depth", reverse=True)
         )
+
+        if self.screen.focused in remove_widgets:
+            self.screen._reset_focus(
+                self.screen.focused,
+                [to_remove for to_remove in remove_widgets if to_remove.can_focus],
+            )
+
         for child in remove_widgets:
             await child._close_messages()
             self._unregister(child)

src/textual/screen.py

@@ -12,6 +12,7 @@
 from ._compositor import Compositor, MapGeometry
 from .timer import Timer
 from ._types import CallbackType
+from .dom import DOMNode
 from .geometry import Offset, Region, Size
 from .reactive import Reactive
 from .renderables.blank import Blank
@@ -224,19 +225,51 @@ def focus_previous(self) -> Widget | None:
         """
         return self._move_focus(-1)
 
-    def _reset_focus(self, widget: Widget) -> None:
+    def _reset_focus(
+        self, widget: Widget, avoiding: list[DOMNode] | None = None
+    ) -> None:
         """Reset the focus when a widget is removed
 
         Args:
             widget (Widget): A widget that is removed.
+            avoiding (list[DOMNode] | None, optional): Optional list of nodes to avoid.
         """
-        if self.focused is widget:
-            for sibling in widget.siblings:
-                if sibling.can_focus:
-                    sibling.focus()
-                    break
-            else:
-                self.focused = None
+
+        avoiding = avoiding or []
+
+        # Make this a NOP if we're being asked to deal with a widget that
+        # isn't actually the currently-focused widget.
+        if self.focused is not widget:
+            return
+
+        # Grab the list of widgets that we can set focus to.
+        focusable_widgets = self.focus_chain
+        if not focusable_widgets:
+            # If there's nothing to focus... give up now.
+            return
+
+        try:
+            # Find the location of the widget we're taking focus from, in
+            # the focus chain.
+            widget_index = focusable_widgets.index(widget)
+        except ValueError:
+            # Seems we can't find it. There's no good reason this should
+            # happen but, on the off-chance, let's go into a "no focus" state.
+            self.set_focus(None)
+            return
+
+        # Now go looking for something before it, that isn't about to be
+        # removed, and which can receive focus, and go focus that.
+        chosen: DOMNode | None = None
+        for candidate in reversed(
+            focusable_widgets[widget_index + 1 :] + focusable_widgets[:widget_index]
+        ):
+            if candidate not in avoiding and candidate.can_focus:
+                chosen = candidate
+                break
+
+        # Go with the what was found.
+        self.set_focus(chosen)
 
     def set_focus(self, widget: Widget | None, scroll_visible: bool = True) -> None:
         """Focus (or un-focus) a widget. A focused widget will receive key events first.

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: