Jonathon Broughton
120 lines
4.6 KiB
Python
120 lines
4.6 KiB
Python
"""Helper module for a speckle object tree flattening."""
|
|
|
|
from collections.abc import Iterable
|
|
from typing import Optional, Tuple, List
|
|
|
|
from specklepy.objects import Base
|
|
from specklepy.objects.other import Instance, Transform
|
|
|
|
|
|
def speckle_print(log_string: str = "banana") -> None:
|
|
|
|
print("\033[92m" + str(log_string) + "\033[0m")
|
|
|
|
|
|
def flatten_base(base: Base) -> Iterable[Base]:
|
|
"""Flatten a base object into an iterable of bases."""
|
|
elements = getattr(base, "elements", getattr(base, "@elements", None))
|
|
if elements is not None:
|
|
for element in elements:
|
|
yield from flatten_base(element)
|
|
yield base
|
|
|
|
|
|
def flatten_base_thorough(base: Base, parent_type: str = None) -> Iterable[Base]:
|
|
"""Take a base and flatten it to an iterable of bases.
|
|
|
|
Args:
|
|
base: The base object to flatten.
|
|
parent_type: The type of the parent object, if any.
|
|
|
|
Yields:
|
|
Base: A flattened base object.
|
|
"""
|
|
if isinstance(base, Base):
|
|
base["parent_type"] = parent_type
|
|
|
|
elements = getattr(base, "elements", getattr(base, "@elements", None))
|
|
if elements:
|
|
try:
|
|
for element in elements:
|
|
# Recursively yield flattened elements of the child
|
|
yield from flatten_base_thorough(element, base.speckle_type)
|
|
except KeyError:
|
|
pass
|
|
elif hasattr(base, "@Lines"):
|
|
categories = base.get_dynamic_member_names()
|
|
|
|
# could be old revit
|
|
try:
|
|
for category in categories:
|
|
print(category)
|
|
if category.startswith("@"):
|
|
category_object: Base = getattr(base, category)[0]
|
|
yield from flatten_base_thorough(
|
|
category_object, category_object.speckle_type
|
|
)
|
|
|
|
except KeyError:
|
|
pass
|
|
|
|
yield base
|
|
|
|
|
|
def extract_base_and_transform(
|
|
base: Base,
|
|
inherited_instance_id: Optional[str] = None,
|
|
transform_list: Optional[List[Transform]] = None,
|
|
) -> Tuple[Base, str, Optional[List[Transform]]]:
|
|
"""
|
|
Traverses Speckle object hierarchies to yield `Base` objects and their transformations.
|
|
Tailored to Speckle's AEC data structures, it covers the newer hierarchical structures
|
|
with Collections and also with patterns found in older Revit specific data.
|
|
|
|
Parameters:
|
|
- base (Base): The starting point `Base` object for traversal.
|
|
- inherited_instance_id (str, optional): The inherited identifier for `Base` objects without a unique ID.
|
|
- transform_list (List[Transform], optional): Accumulated list of transformations from parent to child objects.
|
|
|
|
Yields:
|
|
- tuple: A `Base` object, its identifier, and a list of applicable `Transform` objects or None.
|
|
|
|
The id of the `Base` object is either the inherited identifier for a definition from an instance
|
|
or the one defined in the object.
|
|
"""
|
|
# Derive the identifier for the current `Base` object, defaulting to an inherited one if needed.
|
|
current_id = getattr(base, "id", inherited_instance_id)
|
|
transform_list = transform_list or []
|
|
|
|
if isinstance(base, Instance):
|
|
# Append transformation data and dive into the definition of `Instance` objects.
|
|
if base.transform:
|
|
transform_list.append(base.transform)
|
|
if base.definition:
|
|
yield from extract_base_and_transform(
|
|
base.definition, current_id, transform_list.copy()
|
|
)
|
|
else:
|
|
# Initial yield for the current `Base` object.
|
|
yield base, current_id, transform_list
|
|
|
|
# Process 'elements' and '@elements', typical containers for `Base` objects in AEC models.
|
|
elements_attr = getattr(base, "elements", []) or getattr(base, "@elements", [])
|
|
for element in elements_attr:
|
|
if isinstance(element, Base):
|
|
# Recurse into each `Base` object within 'elements' or '@elements'.
|
|
yield from extract_base_and_transform(
|
|
element, current_id, transform_list.copy()
|
|
)
|
|
|
|
# Recursively process '@'-prefixed properties that are Base objects with 'elements'.
|
|
# This is a common pattern in older Speckle data models, such as those used for Revit commits.
|
|
for attr_name in dir(base):
|
|
if attr_name.startswith("@"):
|
|
attr_value = getattr(base, attr_name)
|
|
# If the attribute is a Base object containing 'elements', recurse into it.
|
|
if isinstance(attr_value, Base) and hasattr(attr_value, "elements"):
|
|
yield from extract_base_and_transform(
|
|
attr_value, current_id, transform_list.copy()
|
|
)
|