refactor(core/tree): handle setting root to None in caller instead of…

… helper
aravinda0 · Mar 6, 2024 · bdd20c4 · bdd20c4
1 parent 63a3946
commit bdd20c4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/qtile_bonsai/core/tree.py b/src/qtile_bonsai/core/tree.py
@@ -385,6 +385,8 @@ def remove(
         br_rm, br_sib, br_rm_nodes = self._remove(
             node, consume_vacant_space=True, normalize=normalize
         )
+        if br_rm is self._root:
+            self._root = None
 
         rm_nodes.extend(br_rm_nodes)
         if br_sib is not None:
@@ -681,7 +683,6 @@ def _remove(
         """
         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