Files
jmwright-CQ-Editor/cq_editor/widgets/object_tree.py
2022-09-16 13:52:44 -05:00

390 lines
13 KiB
Python

from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem, QAction, QMenu, QWidget, QAbstractItemView
from PyQt5.QtCore import Qt, pyqtSlot, pyqtSignal
from pyqtgraph.parametertree import Parameter, ParameterTree
from OCP.AIS import AIS_Line
from OCP.Geom import Geom_Line
from OCP.gp import gp_Dir, gp_Pnt, gp_Ax1
from ..mixins import ComponentMixin
from ..icons import icon
from ..cq_utils import make_AIS, export, to_occ_color, is_obj_empty, get_occ_color
from ..utils import splitter, layout, get_save_filename
class TopTreeItem(QTreeWidgetItem):
def __init__(self,*args,**kwargs):
super(TopTreeItem,self).__init__(*args,**kwargs)
class ObjectTreeItem(QTreeWidgetItem):
props = [{'name': 'Name', 'type': 'str', 'value': ''},
{'name': 'Color', 'type': 'color', 'value': "f4a824"},
{'name': 'Alpha', 'type': 'float', 'value': 0, 'limits': (0,1), 'step': 1e-1},
{'name': 'Visible', 'type': 'bool','value': True}]
def __init__(self,
name,
ais=None,
shape=None,
shape_display=None,
sig=None,
alpha=0.,
color='f4a824',
**kwargs):
super(ObjectTreeItem,self).__init__([name],**kwargs)
self.setFlags( self.flags() | Qt.ItemIsUserCheckable)
self.setCheckState(0,Qt.Checked)
self.ais = ais
self.shape = shape
self.shape_display = shape_display
self.sig = sig
self.properties = Parameter.create(name='Properties',
children=self.props)
self.properties['Name'] = name
self.properties['Alpha'] = ais.Transparency()
self.properties['Color'] = get_occ_color(ais) if ais else color
self.properties.sigTreeStateChanged.connect(self.propertiesChanged)
def propertiesChanged(self,*args):
self.setData(0,0,self.properties['Name'])
self.ais.SetTransparency(self.properties['Alpha'])
self.ais.SetColor(to_occ_color(self.properties['Color']))
self.ais.Redisplay()
if self.properties['Visible']:
self.setCheckState(0,Qt.Checked)
else:
self.setCheckState(0,Qt.Unchecked)
if self.sig:
self.sig.emit()
class CQRootItem(TopTreeItem):
def __init__(self,*args,**kwargs):
super(CQRootItem,self).__init__(['CQ models'],*args,**kwargs)
class HelpersRootItem(TopTreeItem):
def __init__(self,*args,**kwargs):
super(HelpersRootItem,self).__init__(['Helpers'],*args,**kwargs)
class ObjectTree(QWidget,ComponentMixin):
name = 'Object Tree'
_stash = []
preferences = Parameter.create(name='Preferences',children=[
{'name': 'Preserve properties on reload', 'type': 'bool', 'value': False},
{'name': 'Clear all before each run', 'type': 'bool', 'value': True},
{'name': 'STL precision','type': 'float', 'value': .1}])
sigObjectsAdded = pyqtSignal([list],[list,bool])
sigObjectsRemoved = pyqtSignal(list)
sigCQObjectSelected = pyqtSignal(object)
sigAISObjectsSelected = pyqtSignal(list)
sigItemChanged = pyqtSignal(QTreeWidgetItem,int)
sigObjectPropertiesChanged = pyqtSignal()
def __init__(self,parent):
super(ObjectTree,self).__init__(parent)
self.tree = tree = QTreeWidget(self,
selectionMode=QAbstractItemView.ExtendedSelection)
self.properties_editor = ParameterTree(self)
tree.setHeaderHidden(True)
tree.setItemsExpandable(False)
tree.setRootIsDecorated(False)
tree.setContextMenuPolicy(Qt.ActionsContextMenu)
#forward itemChanged singal
tree.itemChanged.connect(\
lambda item,col: self.sigItemChanged.emit(item,col))
#handle visibility changes form tree
tree.itemChanged.connect(self.handleChecked)
self.CQ = CQRootItem()
self.Helpers = HelpersRootItem()
root = tree.invisibleRootItem()
root.addChild(self.CQ)
root.addChild(self.Helpers)
tree.expandToDepth(1)
self._export_STL_action = \
QAction('Export as STL',
self,
enabled=False,
triggered=lambda: \
self.export('stl',
self.preferences['STL precision']))
self._export_STEP_action = \
QAction('Export as STEP',
self,
enabled=False,
triggered=lambda: \
self.export('step'))
self._clear_current_action = QAction(icon('delete'),
'Clear current',
self,
enabled=False,
triggered=self.removeSelected)
self._toolbar_actions = \
[QAction(icon('delete-many'),'Clear all',self,triggered=self.removeObjects),
self._clear_current_action,]
self.prepareMenu()
tree.itemSelectionChanged.connect(self.handleSelection)
tree.customContextMenuRequested.connect(self.showMenu)
self.prepareLayout()
def prepareMenu(self):
self.tree.setContextMenuPolicy(Qt.CustomContextMenu)
self._context_menu = QMenu(self)
self._context_menu.addActions(self._toolbar_actions)
self._context_menu.addActions((self._export_STL_action,
self._export_STEP_action))
def prepareLayout(self):
self._splitter = splitter((self.tree,self.properties_editor),
stretch_factors = (2,1),
orientation=Qt.Vertical)
layout(self,(self._splitter,),top_widget=self)
self._splitter.show()
def showMenu(self,position):
self._context_menu.exec_(self.tree.viewport().mapToGlobal(position))
def menuActions(self):
return {'Tools' : [self._export_STL_action,
self._export_STEP_action]}
def toolbarActions(self):
return self._toolbar_actions
def addLines(self):
origin = (0,0,0)
ais_list = []
for name,color,direction in zip(('X','Y','Z'),
('red','lawngreen','blue'),
((1,0,0),(0,1,0),(0,0,1))):
line_placement = Geom_Line(gp_Ax1(gp_Pnt(*origin),
gp_Dir(*direction)))
line = AIS_Line(line_placement)
line.SetColor(to_occ_color(color))
self.Helpers.addChild(ObjectTreeItem(name,
ais=line))
ais_list.append(line)
self.sigObjectsAdded.emit(ais_list)
def _current_properties(self):
current_params = {}
for i in range(self.CQ.childCount()):
child = self.CQ.child(i)
current_params[child.properties['Name']] = child.properties
return current_params
def _restore_properties(self,obj,properties):
for p in properties[obj.properties['Name']]:
obj.properties[p.name()] = p.value()
@pyqtSlot(dict,bool)
@pyqtSlot(dict)
def addObjects(self,objects,clean=False,root=None):
if root is None:
root = self.CQ
request_fit_view = True if root.childCount() == 0 else False
preserve_props = self.preferences['Preserve properties on reload']
if preserve_props:
current_props = self._current_properties()
if clean or self.preferences['Clear all before each run']:
self.removeObjects()
ais_list = []
#remove empty objects
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)
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)
if request_fit_view:
self.sigObjectsAdded[list,bool].emit(ais_list,True)
else:
self.sigObjectsAdded[list].emit(ais_list)
@pyqtSlot(object,str,object)
def addObject(self,obj,name='',options={}):
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])
@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()]
self.sigObjectsRemoved.emit(removed_items_ais)
@pyqtSlot(bool)
def stashObjects(self,action : bool):
if action:
self._stash = self.CQ.takeChildren()
removed_items_ais = [ch.ais for ch in 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]
self.sigObjectsAdded.emit(ais_list)
@pyqtSlot()
def removeSelected(self):
ixs = self.tree.selectedIndexes()
rows = [ix.row() for ix in ixs]
self.removeObjects(rows)
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
else:
shapes = [item.shape for item in items if item.parent() is self.CQ]
fname = get_save_filename(export_type)
if fname != '':
export(shapes,export_type,fname,precision)
@pyqtSlot()
def handleSelection(self):
items =self.tree.selectedItems()
if len(items) == 0:
self._export_STL_action.setEnabled(False)
self._export_STEP_action.setEnabled(False)
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]
self.sigAISObjectsSelected.emit(ais_objects)
# handle context menu and emit last selected CQ object (if present)
item = items[-1]
if item.parent() is self.CQ:
self._export_STL_action.setEnabled(True)
self._export_STEP_action.setEnabled(True)
self._clear_current_action.setEnabled(True)
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:
self._export_STL_action.setEnabled(True)
self._export_STEP_action.setEnabled(True)
else:
self._export_STL_action.setEnabled(False)
self._export_STEP_action.setEnabled(False)
self._clear_current_action.setEnabled(False)
self.properties_editor.setEnabled(False)
self.properties_editor.clear()
@pyqtSlot(list)
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)
@pyqtSlot(QTreeWidgetItem,int)
def handleChecked(self,item,col):
if type(item) is ObjectTreeItem:
if item.checkState(0):
item.properties['Visible'] = True
else:
item.properties['Visible'] = False