jmwright-CQ-Editor/cq_editor/widgets/object_tree.py

from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem, QAction, QMenu, QWidget, QAbstractItemView
from PyQt5.QtCore import Qt, pyqtSlot, pyqtSignal
from random import randrange as rrr, seed
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, set_color
from .viewer import DEFAULT_FACE_COLOR
from ..utils import splitter, layout, get_save_filename

seed(371353) #preserves colors run to run, needs to be run once globally

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 and ais.HasColor() else get_occ_color(DEFAULT_FACE_COLOR)
        self.properties.sigTreeStateChanged.connect(self.propertiesChanged)

    def propertiesChanged(self, properties, changed):

        changed_prop = changed[0][0]

        self.setData(0,0,self.properties['Name'])
        self.ais.SetTransparency(self.properties['Alpha'])
        
        if changed_prop.name() == 'Color':
            set_color(self.ais, 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

def rand_color(alpha = 0., cfloat=False):
    #helper function to generate a random color dict
    #for CQ-editor's show_object function
    lower = 10
    upper = 100 #not too high to keep color brightness in check
    if cfloat: #for two output types depending on need
        return (
                (rrr(lower,upper)/255),
                (rrr(lower,upper)/255),
                (rrr(lower,upper)/255),
                alpha,
                )
    return {"alpha": alpha,
            "color": (
                      rrr(lower,upper),
                      rrr(lower,upper),
                      rrr(lower,upper),
                     )}