refactor(core/tree): have remove() delegate logic to _remove() so can…

… be used elsewhere

src/qtile_bonsai/core/tree.py

@@ -354,41 +354,43 @@ def _normalize(node: Node, *, recurse: bool):
 
         _normalize(node, recurse=recurse)
 
-    def remove(self, pane: Pane, *, normalize: bool = False) -> Pane | None:
-        removed_nodes = []
-
-        br_remove, br_remove_nodes = self._find_removal_branch(pane)
-        removed_nodes.extend(br_remove_nodes)
-
-        if br_remove is self._root:
-            self._root = None
-            self._notify_subscribers(TreeEvent.node_removed, removed_nodes)
-            return None
+    def remove(
+        self, node: Node, *, normalize: bool = False
+    ) -> tuple[Node, Node | None, Pane | None]:
+        """Remove the provided `node` from the tree.
+
+        Will also notify subscribers of `TreeEvent.node_removed` of all removed nodes.
+
+        Args:
+            `node`:
+                The node to be removed. Note that the resolved removal point could also
+                be an ancestor of this node.
+            `normalize`:
+                Whether removal should lead to all the `node` siblings to re-dimension
+                themselves to take equal amounts of space.
+
+        Returns:
+            3-tuple:
+                1. The `node` or an ancestor node that was the point of removal. This
+                branch is now unlinked from the main tree.
+                2. The sibling node of the removed node. Or `None` if no sibling exists.
+                3. The next pane under the sibling that ought to get focus.
+        """
+        rm_nodes = []
+        next_focus_pane = None
 
-        assert br_remove.parent is not None
-        assert br_remove.operational_sibling is not None
-
-        container = br_remove.parent
-        br_sibling = br_remove.operational_sibling
-        container.children.remove(br_remove)
-
-        if isinstance(container, SplitContainer):
-            # Space redistribution is applicable. Give space to sibling node.
-            br_remove_rect = br_remove.principal_rect
-            br_sibling_rect = br_sibling.principal_rect
-            axis = container.axis
-            start = min(br_remove_rect.coord(axis), br_sibling_rect.coord(axis))
-            br_sibling.transform(
-                axis, start, br_sibling_rect.size(axis) + br_remove_rect.size(axis)
-            )
-            if normalize:
-                self.normalize(container)
+        br_rm, br_sib, br_rm_nodes = self._remove(
+            node, consume_vacant_space=True, normalize=normalize
+        )
 
-        removed_nodes.extend(self._do_post_removal_pruning(br_sibling))
+        rm_nodes.extend(br_rm_nodes)
+        if br_sib is not None:
+            rm_nodes.extend(self._do_post_removal_pruning(br_sib))
+            next_focus_pane = self._pick_mru_pane(self.iter_panes(start=br_sib))
 
-        self._notify_subscribers(TreeEvent.node_removed, removed_nodes)
+        self._notify_subscribers(TreeEvent.node_removed, rm_nodes)
 
-        return self._pick_mru_pane(self.iter_panes(start=br_sibling))
+        return (br_rm, br_sib, next_focus_pane)
 
     def reset(self, from_state: dict | None = None):
         """Clear the current tree. If `from_state` is provided, restore state from
@@ -653,6 +655,40 @@ def walk(node, prefix=""):
 
         return "\n".join(walk(self._root))
 
+    def _remove(
+        self, node: Node, *, consume_vacant_space: bool = True, normalize: bool = False
+    ) -> tuple[Node, Node | None, list[Node]]:
+        """Internal helper for removing node subtrees.
+
+        Returns:
+            A 3-tuple with the following items:
+                1. The resolved node that is the point of removal. May be the provided
+                `node` or an ancestor.
+                2. The sibling of the removed node.
+                3. List of nodes under the removed node.
+        """
+        br_rm, br_rm_nodes = self._find_removal_branch(node)
+        if br_rm is self._root:
+            self._root = None
+            return (br_rm, None, br_rm_nodes)
+
+        container = br_rm.parent
+        br_sib = br_rm.operational_sibling
+        union_rect = br_rm.principal_rect.union(br_sib.principal_rect)
+
+        container.children.remove(br_rm)
+        br_rm.parent = None
+
+        assert br_sib is not None
+
+        if consume_vacant_space and isinstance(container, SplitContainer):
+            br_sib.transform(Axis.x, union_rect.x, union_rect.w)
+            br_sib.transform(Axis.y, union_rect.y, union_rect.h)
+            if normalize:
+                self.normalize(container)
+
+        return (br_rm, br_sib, br_rm_nodes)
+
     def _get_pruning_cases(self) -> list[_PruningCase]:
         return [
             _PruningCase(

src/qtile_bonsai/layout.py

@@ -299,7 +299,7 @@ def remove(self, window: Window) -> Window | None:
         normalize_on_remove = self._tree.get_config(
             "window.normalize_on_remove", level=pane.tab_level
         )
-        next_focus_pane = self._tree.remove(pane, normalize=normalize_on_remove)
+        _, _, next_focus_pane = self._tree.remove(pane, normalize=normalize_on_remove)
         del self._windows_to_panes[window]
         if next_focus_pane is not None:
             return next_focus_pane.window

tests/unit/core/test_tree.py

@@ -1984,7 +1984,7 @@ def test_when_operational_sibling_is_pane_then_it_is_returned_as_next_focus_node
         p1 = tree.tab()
         p2 = tree.split(p1, "x")
 
-        p = tree.remove(p1)
+        _, _, p = tree.remove(p1)
 
         assert p is p2
 
@@ -1998,7 +1998,7 @@ def test_when_operational_sibling_is_sc_then_its_mru_pane_is_returned_as_next_fo
 
         tree.focus(p3)
 
-        p = tree.remove(p1)
+        _, _, p = tree.remove(p1)
 
         assert p is p3
 
@@ -2012,15 +2012,15 @@ def test_when_operational_sibling_is_tc_then_its_mru_pane_is_returned_as_next_fo
 
         tree.focus(p4)
 
-        p = tree.remove(p3)
+        _, _, p = tree.remove(p3)
 
         assert p is p4
 
     def test_when_tree_becomes_empty_then_returns_none_as_nothing_left_for_subsequent_focus(
         self, tree: Tree
     ):
         p1 = tree.tab()
-        p = tree.remove(p1)
+        _, _, p = tree.remove(p1)
 
         assert p is None
 
@@ -2370,6 +2370,8 @@ def test_when_n1_n2_n3_chain_is_sc_sc_tc_then_n1_and_n3_are_linked(
                 p3 = tree.split(p2, "y")
                 tree.tab(p3, new_level=True)
 
+                sc = p2.parent
+
                 tree.remove(p2)
 
                 assert tree_matches_repr(
@@ -2390,7 +2392,7 @@ def test_when_n1_n2_n3_chain_is_sc_sc_tc_then_n1_and_n3_are_linked(
                 )
 
                 assert callback.mock_calls == [
-                    mock.call([p2, p2.parent]),
+                    mock.call([p2, sc]),
                 ]
 
             class TestWhenN1_N2_N3_Is_SC_TC_T:  # noqa: N801