Merge pull request #1704 from zauberzeug/tree-expand-all

Add default_expand_all parameter and warn when using the prop

nicegui/elements/tree.py

@@ -1,12 +1,16 @@
-from typing import Any, Callable, List, Literal, Optional
+from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Set
 
+from typing_extensions import Self
+
+from .. import globals  # pylint: disable=redefined-builtin
 from ..element import Element
 from ..events import GenericEventArguments, ValueChangeEventArguments, handle_event
 
 
 class Tree(Element):
 
-    def __init__(self, nodes: List, *,
+    def __init__(self,
+                 nodes: List[Dict], *,
                  node_key: str = 'id',
                  label_key: str = 'label',
                  children_key: str = 'children',
@@ -31,6 +35,7 @@ def __init__(self, nodes: List, *,
         :param on_expand: callback which is invoked when the node expansion changes
         :param on_tick: callback which is invoked when a node is ticked or unticked
         :param tick_strategy: whether and how to use checkboxes ("leaf", "leaf-filtered" or "strict"; default: ``None``)
+        :param default_expand_all: whether to expand all nodes by default (default: ``False``)
         """
         super().__init__('q-tree')
         self._props['nodes'] = nodes
@@ -62,3 +67,43 @@ def handle_ticked(e: GenericEventArguments) -> None:
             update_prop('ticked', e.args)
             handle_event(on_tick, ValueChangeEventArguments(sender=self, client=self.client, value=e.args))
         self.on('update:ticked', handle_ticked)
+
+    def expand(self, node_keys: Optional[List[str]] = None) -> Self:
+        """Expand the given nodes.
+
+        :param node_keys: list of node keys to expand (default: all nodes)
+        """
+        self._props['expanded'][:] = self._find_node_keys(node_keys).union(self._props['expanded'])
+        self.update()
+        return self
+
+    def collapse(self, node_keys: Optional[List[Dict]] = None) -> Self:
+        """Collapse the given nodes.
+
+        :param node_keys: list of node keys to collapse (default: all nodes)
+        """
+        self._props['expanded'][:] = set(self._props['expanded']).difference(self._find_node_keys(node_keys))
+        self.update()
+        return self
+
+    def _find_node_keys(self, node_keys: Optional[List[str]] = None) -> Set[str]:
+        if node_keys is not None:
+            return set(node_keys)
+
+        CHILDREN_KEY = self._props['children-key']
+        NODE_KEY = self._props['node-key']
+
+        def iterate_nodes(nodes: List[Dict]) -> Iterator[Dict]:
+            for node in nodes:
+                yield node
+                yield from iterate_nodes(node.get(CHILDREN_KEY, []))
+        return {node[NODE_KEY] for node in iterate_nodes(self._props['nodes'])}
+
+    def props(self, add: Optional[str] = None, *, remove: Optional[str] = None) -> Self:
+        super().props(add, remove=remove)
+        if 'default-expand-all' in self._props:
+            # https://github.com/zauberzeug/nicegui/issues/1385
+            del self._props['default-expand-all']
+            globals.log.warning('The prop "default_expand_all" is not supported by `ui.tree`.\n'
+                                'Use ".expand()" instead.')
+        return self

tests/test_tree.py

@@ -0,0 +1,64 @@
+from nicegui import ui
+
+from .screen import Screen
+
+
+def test_tree(screen: Screen):
+    ui.tree([
+        {'id': 'numbers', 'children': [{'id': '1'}, {'id': '2'}]},
+        {'id': 'letters', 'children': [{'id': 'A'}, {'id': 'B'}]},
+    ], label_key='id')
+
+    screen.open('/')
+    screen.should_contain('numbers')
+    screen.should_contain('letters')
+    screen.should_not_contain('1')
+    screen.should_not_contain('2')
+    screen.should_not_contain('A')
+    screen.should_not_contain('B')
+
+    screen.find_by_class('q-icon').click()
+    screen.should_contain('1')
+    screen.should_contain('2')
+
+
+def test_expand_and_collapse_nodes(screen: Screen):
+    tree = ui.tree([
+        {'id': 'numbers', 'children': [{'id': '1'}, {'id': '2'}]},
+        {'id': 'letters', 'children': [{'id': 'A'}, {'id': 'B'}]},
+    ], label_key='id')
+
+    ui.button('Expand all', on_click=tree.expand)
+    ui.button('Collapse all', on_click=tree.collapse)
+    ui.button('Expand numbers', on_click=lambda: tree.expand(['numbers']))
+    ui.button('Collapse numbers', on_click=lambda: tree.collapse(['numbers']))
+
+    screen.open('/')
+    screen.click('Expand all')
+    screen.wait(0.5)
+    screen.should_contain('1')
+    screen.should_contain('2')
+    screen.should_contain('A')
+    screen.should_contain('B')
+
+    screen.click('Collapse all')
+    screen.wait(0.5)
+    screen.should_not_contain('1')
+    screen.should_not_contain('2')
+    screen.should_not_contain('A')
+    screen.should_not_contain('B')
+
+    screen.click('Expand numbers')
+    screen.wait(0.5)
+    screen.should_contain('1')
+    screen.should_contain('2')
+    screen.should_not_contain('A')
+    screen.should_not_contain('B')
+
+    screen.click('Expand all')
+    screen.click('Collapse numbers')
+    screen.wait(0.5)
+    screen.should_not_contain('1')
+    screen.should_not_contain('2')
+    screen.should_contain('A')
+    screen.should_contain('B')

website/more_documentation/tree_documentation.py

@@ -34,26 +34,20 @@ def tree_with_custom_header_and_body():
             <span :props="props">Description: "{{ props.node.description }}"</span>
         ''')
 
-    @text_demo('Expand programmatically', '''
-        The tree can be expanded programmatically by modifying the "expanded" prop.
+    @text_demo('Expand and collapse programmatically', '''
+        The whole tree or individual nodes can be toggled programmatically using the `expand()` and `collapse()` methods.
     ''')
     def expand_programmatically():
-        from typing import List
-
-        def expand(node_ids: List[str]) -> None:
-            t._props['expanded'] = node_ids
-            t.update()
-
         with ui.row():
-            ui.button('all', on_click=lambda: expand(['A', 'B']))
-            ui.button('A', on_click=lambda: expand(['A']))
-            ui.button('B', on_click=lambda: expand(['B']))
-            ui.button('none', on_click=lambda: expand([]))
+            ui.button('+ all', on_click=lambda: t.expand())
+            ui.button('- all', on_click=lambda: t.collapse())
+            ui.button('+ A', on_click=lambda: t.expand(['A']))
+            ui.button('- A', on_click=lambda: t.collapse(['A']))
 
         t = ui.tree([
             {'id': 'A', 'children': [{'id': 'A1'}, {'id': 'A2'}]},
             {'id': 'B', 'children': [{'id': 'B1'}, {'id': 'B2'}]},
-        ], label_key='id')
+        ], label_key='id').expand()
 
     @text_demo('Tree with checkboxes', '''
         The tree can be used with checkboxes by setting the "tick-strategy" prop.