From 6881e7d220609eab671ee1b261b38a6522c87e11 Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Wed, 13 Mar 2024 15:21:46 +0500
Subject: [PATCH] utils.filter_elements docs and tests

---
 .../ifcopenshell/util/selector.py             | 44 +++++++++++++++++--
 .../test/util/test_selector.py                | 15 +++++++
 2 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/src/ifcopenshell-python/ifcopenshell/util/selector.py b/src/ifcopenshell-python/ifcopenshell/util/selector.py
index d252c756bb..198a14fca2 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/selector.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/selector.py
@@ -26,6 +26,7 @@ import ifcopenshell.util.placement
 import ifcopenshell.util.geolocation
 import ifcopenshell.util.classification
 from decimal import Decimal
+from typing import Optional
 
 
 filter_elements_grammar = lark.Lark(
@@ -262,7 +263,44 @@ def get_element_value(element, query):
     return Selector.get_element_value(element, keys)
 
 
-def filter_elements(ifc_file, query, elements=None, edit_in_place=False):
+def filter_elements(
+    ifc_file: ifcopenshell.file,
+    query: str,
+    elements: Optional[set[ifcopenshell.entity_instance]] = None,
+    edit_in_place=False,
+) -> set[ifcopenshell.entity_instance]:
+    """
+    Filter elements based on the provided `query`.
+
+    :param ifc_file: The IFC file object
+    :type ifc_file: ifcopenshell.file.file
+    :param query: Query to execute
+    :type query: str
+    :param elements: Base set of IFC elements for the query.
+        If provided, new elements found for the current query will be added to `elements`.
+        Elements explicitly excluded in the `query` will also be excluded from `elements`
+    :type elements: set[ifcopenshell.entity_instance.entity_instance], optional
+    :param edit_in_place: If `True`, mutate the provided `elements` in place. Defaults to `False`
+    :type edit_in_place: bool
+    :return: Set of filtered elements
+    :rtype: set[ifcopenshell.entity_instance.entity_instance]
+
+    Example:
+
+    .. code:: python
+
+        # Select all walls in the file.
+        elements = ifcopenshell.util.selector.filter_elements(ifc_file, "IfcWall, IfcSlab")
+
+        # Add doors to the elements too.
+        elements = ifcopenshell.util.selector.filter_elements(ifc_file, "IfcDoor", elements)
+
+        # Changed our mind, exclude the slabs.
+        elements = ifcopenshell.util.selector.filter_elements(ifc_file, "! IfcSlab", elements)
+
+        # {#1=IfcWall(...), #2=IfcDoor(...)}
+        print(elements)
+    """
     if elements and not edit_in_place:
         elements = elements.copy()
     transformer = FacetTransformer(ifc_file, elements)
@@ -376,10 +414,10 @@ def set_element_value(ifc_file, element, query, value):
 
 
 class FacetTransformer(lark.Transformer):
-    def __init__(self, ifc_file, elements):
+    def __init__(self, ifc_file: ifcopenshell.file, elements: Optional[set[ifcopenshell.entity_instance]] = None):
         self.file = ifc_file
         self.results = []
-        self.elements = elements or set()
+        self.elements = set() if elements is None else elements
         self.container_parents = {}
         self.container_trees = {}
 
diff --git a/src/ifcopenshell-python/test/util/test_selector.py b/src/ifcopenshell-python/test/util/test_selector.py
index b02ab7e811..ff312bec61 100644
--- a/src/ifcopenshell-python/test/util/test_selector.py
+++ b/src/ifcopenshell-python/test/util/test_selector.py
@@ -251,6 +251,21 @@ class TestFilterElements(test.bootstrap.IFC4):
         assert subject.filter_elements(self.file, "IfcWall, Name=Foo + IfcSlab") == {element, element2}
         assert subject.filter_elements(self.file, "IfcWall, Name=Foo + IfcSlab, Name=Bar") == {element, element2}
 
+    def test_using_elements_argument(self):
+        wall = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWall")
+        slab = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcSlab")
+
+        # keep elements unaffected by expression
+        assert subject.filter_elements(self.file, "IfcWall", {slab}) == {wall, slab}
+
+        # filter out excluded elements
+        assert subject.filter_elements(self.file, "IfcWall, ! IfcSlab", {slab}) == {wall}
+
+        # edit_in_place to update original set
+        original_set = set()
+        new_set = subject.filter_elements(self.file, "IfcWall", original_set, edit_in_place=True)
+        assert new_set == original_set
+
 
 class TestSelector(test.bootstrap.IFC4):
     def test_selecting_by_class(self):