Route NestedTree's collapse calls through one loop

tests/test_nested.py

@@ -161,6 +161,17 @@ def test_collapse_all_reaches_hidden_subtrees(self):
         tree.expand_all()
         self.assertEqual(len(list(tree.positions())), 4)
 
+    def test_collapse_deeply_nested(self):
+        # no recursion per NestedTree wrapped directly in another
+        tree, x = self.inner, (0,)
+        for _ in range(1100):
+            tree, x = NestedTree(tree), (x,)
+        self.assertTrue(tree.collapsible(x))
+        tree.collapse(x)
+        self.assertTrue(self.inner.is_collapsed((0,)))
+        tree.expand(x)
+        self.assertFalse(tree.is_collapsed(x))
+
     def test_collapse_all_wide(self):
         outer = CollapsibleTree(SimpleTree([(Node('R'), [
             (Node(str(i)), [(Node('c'), None)]) for i in range(1500)])]))

urwidtrees/nested.py

@@ -98,43 +98,38 @@ def decorate(self, pos, widget, is_first=True):
         return self._get_decorated_entry(pos, widget, is_first)
 
     # Collapse API
-    def _get_subtree_for(self, pos):
-        """returns Tree that manages pos[-1]"""
-        return self._trees(pos)[-1]
+    def _on_subtree(self, pos, method, default=None):
+        """
+        call `method` with pos[-1] on the tree that manages it, if that tree
+        can collapse; otherwise return `default`. A NestedTree nested
+        directly in another hands on to its own managing tree.
+        """
+        tree = self
+        while isinstance(tree, NestedTree):
+            tree, pos = tree._trees(pos)[-1], pos[-1]
+        if isinstance(tree, CollapseMixin):
+            return getattr(tree, method)(pos)
+        return default
 
     def collapsible(self, pos):
-        res = False
-        subtree = self._get_subtree_for(pos)
-        if isinstance(subtree, (CollapseMixin, NestedTree)):
-            res = subtree.collapsible(pos[-1])
-        return res
+        return self._on_subtree(pos, 'collapsible', False)
 
     def is_collapsed(self, pos):
-        res = False
-        subtree = self._get_subtree_for(pos)
-        if isinstance(subtree, (CollapseMixin, NestedTree)):
-            res = subtree.is_collapsed(pos[-1])
-        return res
+        return self._on_subtree(pos, 'is_collapsed', False)
 
     def toggle_collapsed(self, pos):
-        subtree = self._get_subtree_for(pos)
-        if isinstance(subtree, (CollapseMixin, NestedTree)):
-            subtree.toggle_collapsed(pos[-1])
+        self._on_subtree(pos, 'toggle_collapsed')
 
     def collapse(self, pos):
-        subtree = self._get_subtree_for(pos)
-        if isinstance(subtree, (CollapseMixin, NestedTree)):
-            subtree.collapse(pos[-1])
+        self._on_subtree(pos, 'collapse')
 
     def collapse_all(self):
         for tree in self._expand_all():
             if isinstance(tree, (CollapseMixin, NestedTree)):
                 tree.collapse_all()
 
     def expand(self, pos):
-        subtree = self._get_subtree_for(pos)
-        if isinstance(subtree, (CollapseMixin, NestedTree)):
-            subtree.expand(pos[-1])
+        self._on_subtree(pos, 'expand')
 
     def expand_all(self):
         self._expand_all()