diff --git a/src/blenderbim/blenderbim/tool/search.py b/src/blenderbim/blenderbim/tool/search.py index 834b99304c..7e11e4ab87 100644 --- a/src/blenderbim/blenderbim/tool/search.py +++ b/src/blenderbim/blenderbim/tool/search.py @@ -5,15 +5,17 @@ import blenderbim.core.tool import blenderbim.tool as tool import ifcopenshell.util.selector from ifcopenshell.util.selector import Selector +from blenderbim.bim.prop import BIMFacet +from typing import Union, Literal class Search(blenderbim.core.tool.Search): @classmethod - def get_group_query(cls, group): + def get_group_query(cls, group: ifcopenshell.entity_instance) -> str: return json.loads(group.Description)["query"] @classmethod - def get_filter_groups(cls, module): + def get_filter_groups(cls, module: str) -> bpy.types.bpy_prop_collection: if module == "search": return bpy.context.scene.BIMSearchProperties.filter_groups elif module == "csv": @@ -24,13 +26,13 @@ class Search(blenderbim.core.tool.Search): return bpy.context.active_object.data.BIMCameraProperties.exclude_filter_groups @classmethod - def import_filter_query(cls, query, filter_groups): + def import_filter_query(cls, query: str, filter_groups: bpy.types.bpy_prop_collection) -> None: filter_groups.clear() transformer = ImportFilterQueryTransformer(filter_groups) transformer.transform(ifcopenshell.util.selector.filter_elements_grammar.parse(query)) @classmethod - def export_filter_query(cls, filter_groups): + def export_filter_query(cls, filter_groups: bpy.types.bpy_prop_collection) -> str: query = [] for filter_group in filter_groups: filter_group_query = [] @@ -87,13 +89,15 @@ class Search(blenderbim.core.tool.Search): return " + ".join(query) @classmethod - def get_comparison_and_value(cls, ifc_filter): + def get_comparison_and_value( + cls, ifc_filter: BIMFacet + ) -> Union[tuple[Literal["!="], str], tuple[Literal["="], str]]: if ifc_filter.value.startswith("!="): return ("!=", cls.wrap_value(ifc_filter, ifc_filter.value[2:].strip())) return ("=", cls.wrap_value(ifc_filter, ifc_filter.value.strip())) @classmethod - def wrap_value(cls, ifc_filter, value): + def wrap_value(cls, ifc_filter: BIMFacet, value: str) -> str: if value.startswith("/") and value.endswith("/"): return value elif value in ("NULL", "TRUE", "FALSE"): @@ -101,7 +105,7 @@ class Search(blenderbim.core.tool.Search): return '"' + value.replace('"', '\\"') + '"' @classmethod - def from_selector_query(cls, query): + def from_selector_query(cls, query: str) -> list[ifcopenshell.entity_instance]: """Returns a list of products from a selector query""" return Selector().parse(tool.Ifc.get(), query) diff --git a/src/ifcopenshell-python/ifcopenshell/util/selector.py b/src/ifcopenshell-python/ifcopenshell/util/selector.py index 2ac14f0c3f..9c2236b2d3 100644 --- a/src/ifcopenshell-python/ifcopenshell/util/selector.py +++ b/src/ifcopenshell-python/ifcopenshell/util/selector.py @@ -711,7 +711,9 @@ class FacetTransformer(lark.Transformer): class Selector: @classmethod - def parse(cls, ifc_file, query, elements=None): + def parse( + cls, ifc_file: ifcopenshell.file, query: str, elements: Optional[list[ifcopenshell.entity_instance]] = None + ) -> list[ifcopenshell.entity_instance]: cls.file = ifc_file cls.elements = elements l = lark.Lark( diff --git a/src/ifcopenshell-python/test/util/test_selector.py b/src/ifcopenshell-python/test/util/test_selector.py index d913313954..9a5f9cadce 100644 --- a/src/ifcopenshell-python/test/util/test_selector.py +++ b/src/ifcopenshell-python/test/util/test_selector.py @@ -21,6 +21,7 @@ import test.bootstrap import ifcopenshell.api import ifcopenshell.util.selector as subject import ifcopenshell.util.placement +import ifcopenshell.util.pset import numpy as np @@ -293,6 +294,10 @@ class TestSetElementValue(test.bootstrap.IFC4): class TestSelector(test.bootstrap.IFC4): + def test_selecting_from_specified_elements(self): + elements = [ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWall") for _ in range(2)] + assert subject.Selector.parse(self.file, ".IfcWall", elements[:1]) == [elements[0]] + def test_selecting_by_class(self): element = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWall") ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcSlab")