typing

src/ifcopenshell-python/ifcopenshell/util/attribute.py

@@ -17,13 +17,22 @@
 # along with IfcOpenShell.  If not, see <http://www.gnu.org/licenses/>.
 
 import ifcopenshell.ifcopenshell_wrapper as ifcopenshell_wrapper
-from typing import Union
+from typing import Union, Literal
+
+PrimitiveType = Literal["entity", "string", "float", "integer", "boolean", "enum", "binary"]
+ComplexPrimitiveType = Literal["list", "array", "set"]
+PrimitiveTypeOutput = Union[
+    PrimitiveType,
+    tuple[ComplexPrimitiveType, "PrimitiveTypeOutput"],
+    tuple[Literal["select"], tuple["PrimitiveTypeOutput", ...]],
+    None,
+]
 
 
 def get_primitive_type(
     attribute_or_data_type: Union[ifcopenshell_wrapper.attribute, ifcopenshell_wrapper.parameter_type],
-) -> Union[str, tuple[str, list[str]]]:
-    if hasattr(attribute_or_data_type, "type_of_attribute"):
+) -> PrimitiveTypeOutput:
+    if isinstance(attribute_or_data_type, ifcopenshell_wrapper.attribute):
         data_type = str(attribute_or_data_type.type_of_attribute())
     else:
         data_type = str(attribute_or_data_type)
@@ -56,8 +65,16 @@ def get_primitive_type(
 
 
 def get_enum_items(attribute: ifcopenshell_wrapper.attribute) -> tuple[str, ...]:
-    return attribute.type_of_attribute().declared_type().enumeration_items()
+    named_type = attribute.type_of_attribute().as_named_type()
+    assert named_type
+    enumeration = named_type.declared_type().as_enumeration_type()
+    assert enumeration
+    return enumeration.enumeration_items()
 
 
-def get_select_items(attribute: ifcopenshell_wrapper.attribute) -> tuple[ifcopenshell_wrapper.entity, ...]:
-    return attribute.type_of_attribute().declared_type().select_list()
+def get_select_items(attribute: ifcopenshell_wrapper.attribute) -> tuple[ifcopenshell_wrapper.declaration, ...]:
+    named_type = attribute.type_of_attribute().as_named_type()
+    assert named_type
+    select_type = named_type.declared_type().as_select_type()
+    assert select_type
+    return select_type.select_list()