diff --git a/cq_editor/widgets/object_tree.py b/cq_editor/widgets/object_tree.py index d6651615..1e216e73 100644 --- a/cq_editor/widgets/object_tree.py +++ b/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): diff --git a/cq_editor/widgets/viewer.py b/cq_editor/widgets/viewer.py index 05cdad12..c3bce2cd 100644 --- a/cq_editor/widgets/viewer.py +++ b/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)