Better Assembly Handling

cq_editor/widgets/object_tree.py

@@ -1,3 +1,4 @@
+from cadquery import Location, Assembly
 from PyQt5.QtWidgets import (
     QTreeWidget,
     QTreeWidgetItem,
@@ -24,6 +25,8 @@
     is_obj_empty,
     get_occ_color,
     set_color,
+    set_transparency,
+    to_compound,
 )
 from .viewer import DEFAULT_FACE_COLOR
 from ..utils import splitter, layout, get_save_filename
@@ -129,6 +132,7 @@ class ObjectTree(QWidget, ComponentMixin):
         children=[
             {"name": "Preserve properties on reload", "type": "bool", "value": False},
             {"name": "Clear all before each run", "type": "bool", "value": True},
+            {"name": "Merge Assemblies", "type": "bool", "value": False},
             {"name": "STL precision", "type": "float", "value": 0.1},
         ],
     )
@@ -151,7 +155,7 @@ def __init__(self, parent):
         self.properties_editor = ParameterTree(self)
 
         tree.setHeaderHidden(True)
-        tree.setItemsExpandable(False)
+        tree.setItemsExpandable(True)
         tree.setRootIsDecorated(False)
         tree.setContextMenuPolicy(Qt.ActionsContextMenu)
 
@@ -287,19 +291,121 @@ def addLines(self):
 
         self.sigObjectsAdded.emit(ais_list)
 
-    def _current_properties(self):
+    def _item_path(self, item):
+        """Stable identity of an item across runs: names from the CQ root down."""
+        parts = []
+        node = item
+        while node is not None and node is not self.CQ:
+            parts.append(node.properties["Name"])
+            node = node.parent()
+        return "/".join(reversed(parts))
 
+    def _current_properties(self):
+        """
+        Snapshot every CQ item's properties before a reload, keyed by tree
+        path so nested parts that share a name do not collide.
+        """
         current_params = {}
         for i in range(self.CQ.childCount()):
-            child = self.CQ.child(i)
-            current_params[child.properties["Name"]] = child.properties
+            for it in self._iter_subtree(self.CQ.child(i)):
+                current_params[self._item_path(it)] = it.properties
 
         return current_params
 
     def _restore_properties(self, obj, properties):
+        """
+        Re-apply a snapshot from _current_properties to obj and all its
+        descendants, matching items by tree path.
+        """
+        for it in self._iter_subtree(obj):
+            key = self._item_path(it)
+            if key in properties:
+                for p in properties[key]:
+                    it.properties[p.name()] = p.value()
+
+    def _build_assembly_item(
+        self, node, label, options, parent_loc, inherited_color, ais_list
+    ):
+        """
+        Recursively build the tree item for one assembly node, mirroring
+        the hierarchy. Accumulates world location and nearest-ancestor color.
+        """
+        world = parent_loc * node.loc
+        color = node.color if node.color is not None else inherited_color
+
+        ais = None
+        shape = None
+        if node.obj is not None:
+            # A node can have both a shape and children
+            shape = to_compound(node.obj).moved(world)
+            ais, _ = make_AIS(shape, options)
+            if color is not None:
+                r, g, b, a = color.toTuple()
+                set_color(ais, to_occ_color((r, g, b)))
+                set_transparency(ais, a)
+            ais_list.append(ais)
+
+        item = ObjectTreeItem(
+            label, shape=shape, ais=ais, sig=self.sigObjectPropertiesChanged
+        )
 
-        for p in properties[obj.properties["Name"]]:
-            obj.properties[p.name()] = p.value()
+        for child in node.children:
+            item.addChild(
+                self._build_assembly_item(
+                    child, child.name, options, world, color, ais_list
+                )
+            )
+
+        if node.children:
+            item.setFlags(item.flags() | Qt.ItemIsAutoTristate)
+
+        return item
+
+    def _build_items(self, name, shape, options):
+        """
+        Build the ObjectTreeItem(s) for one shown object. Assemblies explode
+        into one item per part. Everything else is one item.
+        """
+        # Explode assemblies into per-part items
+        if isinstance(shape, Assembly) and not self.preferences["Merge Assemblies"]:
+            ais_list = []
+            item = self._build_assembly_item(
+                shape, name, options, Location(), None, ais_list
+            )
+            return [item], ais_list
+
+        ais, shape_display = make_AIS(shape, options)
+        item = ObjectTreeItem(
+            name,
+            shape=shape,
+            shape_display=shape_display,
+            ais=ais,
+            sig=self.sigObjectPropertiesChanged,
+        )
+        return [item], [ais]
+
+    def _iter_subtree(self, item):
+        """Yield item and every descendant."""
+        yield item
+        for i in range(item.childCount()):
+            yield from self._iter_subtree(item.child(i))
+
+    def _under_cq(self, item):
+        """True if item is anywhere beneath the CQ root (exluding Helpers)."""
+        p = item.parent()
+        while p is not None:
+            if p is self.CQ:
+                return True
+            p = p.parent()
+        return False
+
+    def _subtree_ais(self, tops):
+        return [
+            it.ais
+            for top in tops
+            for it in self._iter_subtree(top)
+            if it.ais is not None
+        ]
 
     @pyqtSlot(dict, bool)
     @pyqtSlot(dict)
@@ -323,23 +429,14 @@ def addObjects(self, objects, clean=False, root=None):
         objects_f = {k: v for k, v in objects.items() if not is_obj_empty(v.shape)}
 
         for name, obj in objects_f.items():
-            ais, shape_display = make_AIS(obj.shape, obj.options)
-
-            child = ObjectTreeItem(
-                name,
-                shape=obj.shape,
-                shape_display=shape_display,
-                ais=ais,
-                sig=self.sigObjectPropertiesChanged,
-            )
+            top_items, obj_ais = self._build_items(name, obj.shape, obj.options)
+            for item in top_items:
+                if preserve_props and name in current_props:
+                    self._restore_properties(item, current_props)
+                self.CQ.addChild(item)
+                self.tree.expandItem(item)
 
-            if preserve_props and name in current_props:
-                self._restore_properties(child, current_props)
-
-            if child.properties["Visible"]:
-                ais_list.append(ais)
-
-            root.addChild(child)
+            ais_list.extend(obj_ais)
 
         if request_fit_view:
             self.sigObjectsAdded[list, bool].emit(ais_list, True)
@@ -356,28 +453,26 @@ def addObject(self, obj, name="", options=None):
 
         root = self.CQ
 
-        ais, shape_display = make_AIS(obj, options)
-
-        root.addChild(
-            ObjectTreeItem(
-                name,
-                shape=obj,
-                shape_display=shape_display,
-                ais=ais,
-                sig=self.sigObjectPropertiesChanged,
-            )
-        )
-
-        self.sigObjectsAdded.emit([ais])
+        top_items, ais_list = self._build_items(name, obj, options)
+        for item in top_items:
+            self.CQ.addChild(item)
+        self.sigObjectsAdded.emit(ais_list)
 
     @pyqtSlot(list)
     @pyqtSlot()
     def removeObjects(self, objects=None):
 
-        if objects:
-            removed_items_ais = [self.CQ.takeChild(i).ais for i in objects]
-        else:
-            removed_items_ais = [ch.ais for ch in self.CQ.takeChildren()]
+        taken = (
+            [self.CQ.takeChild(i) for i in objects]
+            if objects
+            else self.CQ.takeChildren()
+        )
+        removed_items_ais = [
+            it.ais
+            for top in taken
+            for it in self._iter_subtree(top)
+            if it.ais is not None
+        ]
 
         self.sigObjectsRemoved.emit(removed_items_ais)
 
@@ -386,33 +481,41 @@ def stashObjects(self, action: bool):
 
         if action:
             self._stash = self.CQ.takeChildren()
-            removed_items_ais = [ch.ais for ch in self._stash]
+            # removed_items_ais = [ch.ais for ch in self._stash]
+            removed_items_ais = self._subtree_ais(self._stash)
             self.sigObjectsRemoved.emit(removed_items_ais)
         else:
             self.removeObjects()
             self.CQ.addChildren(self._stash)
-            ais_list = [el.ais for el in self._stash]
+            ais_list = self._subtree_ais(self._stash)
             self.sigObjectsAdded.emit(ais_list)
 
     @pyqtSlot()
     def removeSelected(self):
-
-        ixs = self.tree.selectedIndexes()
-        rows = [ix.row() for ix in ixs]
-
-        self.removeObjects(rows)
+        tops = [it for it in self.tree.selectedItems() if it.parent() is self.CQ]
+        removed_items_ais = self._subtree_ais(tops)
+        for it in tops:
+            self.CQ.removeChild(it)
+        self.sigObjectsRemoved.emit(removed_items_ais)
 
     def export(self, export_type, precision=None):
 
         items = self.tree.selectedItems()
 
-        # if CQ models is selected get all children
-        if [item for item in items if item is self.CQ]:
-            CQ = self.CQ
-            shapes = [CQ.child(i).shape for i in range(CQ.childCount())]
-        # otherwise collect all selected children of CQ
+        # If the CQ root is selected, take all top-level objects
+        if any(it is self.CQ for it in items):
+            roots = [self.CQ.child(i) for i in range(self.CQ.childCount())]
+        # Otherwise take every selected item anywhere under CQ
         else:
-            shapes = [item.shape for item in items if item.parent() is self.CQ]
+            roots = [it for it in items if self._under_cq(it)]
+
+        seen = set()
+        shapes = []
+        for r in roots:
+            for it in self._iter_subtree(r):
+                if it.shape is not None and id(it) not in seen:
+                    seen.add(id(it))
+                    shapes.append(it.shape)
 
         fname = get_save_filename(export_type)
         if fname != "":
@@ -428,16 +531,24 @@ def handleSelection(self):
             return
 
         # emit list of all selected ais objects (might be empty)
-        ais_objects = [item.ais for item in items if item.parent() is self.CQ]
+        # ais_objects = [item.ais for item in items if item.parent() is self.CQ]
+        ais_objects = [
+            it.ais
+            for sel in items
+            if self._under_cq(sel)
+            for it in self._iter_subtree(sel)
+            if it.ais is not None
+        ]
         self.sigAISObjectsSelected.emit(ais_objects)
 
-        # handle context menu and emit last selected CQ  object (if present)
+        # Context menu + last-selected object
         item = items[-1]
-        if item.parent() is self.CQ:
+        if self._under_cq(item):
             self._export_STL_action.setEnabled(True)
             self._export_STEP_action.setEnabled(True)
             self._clear_current_action.setEnabled(True)
-            self.sigCQObjectSelected.emit(item.shape)
+            if item.shape is not None:
+                self.sigCQObjectSelected.emit(item.shape)
             self.properties_editor.setParameters(item.properties, showTop=False)
             self.properties_editor.setEnabled(True)
         elif item is self.CQ and item.childCount() > 0:
@@ -455,12 +566,12 @@ def handleGraphicalSelection(self, shapes):
 
         self.tree.clearSelection()
 
-        CQ = self.CQ
-        for i in range(CQ.childCount()):
-            item = CQ.child(i)
-            for shape in shapes:
-                if item.ais.Shape().IsEqual(shape):
-                    item.setSelected(True)
+        for item in self._iter_subtree(self.CQ):
+            ais = getattr(item, "ais", None)
+            if ais is None:
+                continue
+            if any(ais.Shape().IsEqual(shape) for shape in shapes):
+                item.setSelected(True)
 
     @pyqtSlot(QTreeWidgetItem, int)
     def handleChecked(self, item, col):

cq_editor/widgets/viewer.py

@@ -318,6 +318,9 @@ def display_many(self, ais_list, fit=None):
     @pyqtSlot(QTreeWidgetItem, int)
     def update_item(self, item, col):
 
+        if getattr(item, "ais", None) is None:
+            return
+
         ctx = self._get_context()
         if item.checkState(0):
             ctx.Display(item.ais, True)