NLSA
06b66145b6
build and deploy Speckle functions / publish-automate-function-version (push) Has been cancelled
88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
# =============================================================================
|
|
# traversal.py
|
|
# Walks the nested Speckle Collection tree generically.
|
|
#
|
|
# Expected structure:
|
|
# Root Collection
|
|
# └── Collection
|
|
# └── Collection
|
|
# └── Object (leaf BIM element)
|
|
#
|
|
# Collections can nest to any depth. Every non-Collection leaf is yielded.
|
|
# =============================================================================
|
|
|
|
from typing import Generator
|
|
from specklepy.objects.base import Base
|
|
|
|
|
|
def is_collection(obj) -> bool:
|
|
"""Returns True if this object is a Speckle Collection node (not a leaf element)."""
|
|
speckle_type = getattr(obj, "speckle_type", "") or ""
|
|
return "Collection" in speckle_type
|
|
|
|
|
|
def get_children(obj) -> list:
|
|
"""
|
|
Safely get the 'elements' list from a Base/Collection object.
|
|
Handles 'elements', '@elements', and '_elements' variants.
|
|
"""
|
|
for key in ["elements", "@elements", "_elements"]:
|
|
try:
|
|
val = obj[key]
|
|
if val is not None:
|
|
return list(val)
|
|
except Exception:
|
|
continue
|
|
return []
|
|
|
|
|
|
def traverse(root: Base) -> Generator[Base, None, None]:
|
|
"""
|
|
Walk the full Speckle object tree from the root Base object.
|
|
Yields every non-Collection leaf object found at any depth.
|
|
"""
|
|
yield from _walk(root)
|
|
|
|
|
|
def _walk(obj):
|
|
"""Recursively walk: descend into Collections, yield leaf objects."""
|
|
if obj is None:
|
|
return
|
|
|
|
children = get_children(obj)
|
|
|
|
if is_collection(obj):
|
|
for child in children:
|
|
yield from _walk(child)
|
|
else:
|
|
# Leaf object — yield it
|
|
yield obj
|
|
# Also check for nested children (e.g. curtain wall sub-elements)
|
|
for child in children:
|
|
if child is not None and not is_collection(child):
|
|
yield from _walk(child)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Debug helper
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
def print_tree(obj: Base, indent: int = 0, max_depth: int = 5):
|
|
"""Print the object tree structure for debugging."""
|
|
if indent > max_depth:
|
|
return
|
|
|
|
prefix = " " * indent
|
|
name = getattr(obj, "name", None) or ""
|
|
speckle_type = getattr(obj, "speckle_type", "") or ""
|
|
children = get_children(obj)
|
|
child_count = f" ({len(children)} children)" if children else ""
|
|
|
|
print(f"{prefix}├─ [{speckle_type}] name={name!r}{child_count}")
|
|
|
|
for child in children[:5]:
|
|
print_tree(child, indent + 1, max_depth)
|
|
|
|
if len(children) > 5:
|
|
print(f"{prefix} ... and {len(children) - 5} more")
|