From ff3e714a38c4447759d140959055299d7a2fc32a Mon Sep 17 00:00:00 2001
From: Falko Schindler <falko@zauberzeug.com>
Date: Mon, 12 Aug 2024 13:41:52 +0200
Subject: [PATCH] allow filtering nodes

---
 nicegui/elements/tree.py                        |  6 +++---
 tests/test_tree.py                              | 17 +++++++++++++++++
 .../documentation/content/tree_documentation.py | 12 ++++++++++++
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/nicegui/elements/tree.py b/nicegui/elements/tree.py
index 01010bc33..e814b7bc8 100644
--- a/nicegui/elements/tree.py
+++ b/nicegui/elements/tree.py
@@ -3,11 +3,11 @@
 from typing_extensions import Self
 
 from .. import helpers
-from ..element import Element
 from ..events import GenericEventArguments, ValueChangeEventArguments, handle_event
+from .mixins.filter_element import FilterElement
 
 
-class Tree(Element):
+class Tree(FilterElement):
 
     def __init__(self,
                  nodes: List[Dict], *,
@@ -36,7 +36,7 @@ def __init__(self,
         :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``)
         """
-        super().__init__('q-tree')
+        super().__init__(tag='q-tree', filter=None)
         self._props['nodes'] = nodes
         self._props['node-key'] = node_key
         self._props['label-key'] = label_key
diff --git a/tests/test_tree.py b/tests/test_tree.py
index bebb56e2e..b59c44c42 100644
--- a/tests/test_tree.py
+++ b/tests/test_tree.py
@@ -108,3 +108,20 @@ def test_tick_untick_node_or_nodes(screen: Screen):
 
     screen.click('Untick all')
     screen.should_contain('Ticked: []')
+
+
+def test_filter(screen: Screen):
+    t = ui.tree([
+        {'id': 'fruits', 'children': [{'id': 'Apple'}, {'id': 'Banana'}, {'id': 'Cherry'}]},
+    ], label_key='id', tick_strategy='leaf-filtered').expand()
+    ui.button('Filter', on_click=lambda: t.set_filter('a'))
+
+    screen.open('/')
+    screen.should_contain('Apple')
+    screen.should_contain('Banana')
+    screen.should_contain('Cherry')
+
+    screen.click('Filter')
+    screen.should_contain('Apple')
+    screen.should_contain('Banana')
+    screen.should_not_contain('Cherry')
diff --git a/website/documentation/content/tree_documentation.py b/website/documentation/content/tree_documentation.py
index c0ae0e4fb..893b4c485 100644
--- a/website/documentation/content/tree_documentation.py
+++ b/website/documentation/content/tree_documentation.py
@@ -94,4 +94,16 @@ def tick_programmatically():
         ui.button('Untick all', on_click=t.untick)
 
 
+@doc.demo('Filter nodes', '''
+    You can filter nodes by setting the `filter` property.
+    The tree will only show nodes that match the filter.
+''')
+def filter_nodes():
+    t = ui.tree([
+        {'id': 'fruits', 'children': [{'id': 'Apple'}, {'id': 'Banana'}]},
+        {'id': 'vegetables', 'children': [{'id': 'Potato'}, {'id': 'Tomato'}]},
+    ], label_key='id').expand()
+    ui.input('filter').bind_value_to(t, 'filter')
+
+
 doc.reference(ui.tree)