From c0fecb160b8dc9f5b2052f2c714bdb66819fdc70 Mon Sep 17 00:00:00 2001 From: Jonathon Broughton Date: Thu, 16 May 2024 18:45:32 +0200 Subject: [PATCH] Enhance robustness of the basic flatten command (#20) Older connector versions and still some edge cases in the current set refer to elements with the @ decorator (Groups in Revit for example) --- flatten.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/flatten.py b/flatten.py index c180564..574c118 100644 --- a/flatten.py +++ b/flatten.py @@ -6,8 +6,22 @@ from specklepy.objects import Base def flatten_base(base: Base) -> Iterable[Base]: - """Take a base and flatten it to an iterable of bases.""" - if hasattr(base, "elements") and base.elements is not None: - for element in base["elements"]: + """Flatten a base object into an iterable of bases. + + This function recursively traverses the `elements` or `@elements` attribute of the + base object, yielding each nested base object. + + Args: + base (Base): The base object to flatten. + + Yields: + Base: Each nested base object in the hierarchy. + """ + # Attempt to get the elements attribute, fallback to @elements if necessary + elements = getattr(base, "elements", getattr(base, "@elements", None)) + + if elements is not None: + for element in elements: yield from flatten_base(element) + yield base