diff --git a/cq_editor/widgets/debugger.py b/cq_editor/widgets/debugger.py index c56c0bf..120d11f 100644 --- a/cq_editor/widgets/debugger.py +++ b/cq_editor/widgets/debugger.py @@ -6,9 +6,15 @@ from typing import List import cadquery as cq from PyQt5 import QtCore -from PyQt5.QtCore import Qt, QObject, pyqtSlot, pyqtSignal, QEventLoop, QAbstractTableModel -from PyQt5.QtWidgets import (QAction, - QTableView) +from PyQt5.QtCore import ( + Qt, + QObject, + pyqtSlot, + pyqtSignal, + QEventLoop, + QAbstractTableModel, +) +from PyQt5.QtWidgets import QAction, QTableView from logbook import info from path import Path from pyqtgraph.parametertree import Parameter @@ -18,46 +24,43 @@ from random import randrange as rrr, seed from ..cq_utils import find_cq_objects, reload_cq from ..mixins import ComponentMixin -DUMMY_FILE = '' +DUMMY_FILE = "" class DbgState(Enum): - STEP = auto() CONT = auto() STEP_IN = auto() RETURN = auto() -class DbgEevent(object): - LINE = 'line' - CALL = 'call' - RETURN = 'return' +class DbgEevent(object): + LINE = "line" + CALL = "call" + RETURN = "return" + class LocalsModel(QAbstractTableModel): + HEADER = ("Name", "Type", "Value") - HEADER = ('Name','Type', 'Value') - - def __init__(self,parent): - - super(LocalsModel,self).__init__(parent) + def __init__(self, parent): + super(LocalsModel, self).__init__(parent) self.frame = None - def update_frame(self,frame): - - self.frame = \ - [(k,type(v).__name__, str(v)) for k,v in frame.items() if not k.startswith('_')] - - - def rowCount(self,parent=QtCore.QModelIndex()): + def update_frame(self, frame): + self.frame = [ + (k, type(v).__name__, str(v)) + for k, v in frame.items() + if not k.startswith("_") + ] + def rowCount(self, parent=QtCore.QModelIndex()): if self.frame: return len(self.frame) else: return 0 - def columnCount(self,parent=QtCore.QModelIndex()): - + def columnCount(self, parent=QtCore.QModelIndex()): return 3 def headerData(self, section, orientation, role=Qt.DisplayRole): @@ -74,13 +77,11 @@ class LocalsModel(QAbstractTableModel): return QtCore.QVariant() -class LocalsView(QTableView,ComponentMixin): +class LocalsView(QTableView, ComponentMixin): + name = "Variables" - name = 'Variables' - - def __init__(self,parent): - - super(LocalsView,self).__init__(parent) + def __init__(self, parent): + super(LocalsView, self).__init__(parent) ComponentMixin.__init__(self) header = self.horizontalHeader() @@ -90,193 +91,225 @@ class LocalsView(QTableView,ComponentMixin): vheader.setVisible(False) @pyqtSlot(dict) - def update_frame(self,frame): - + def update_frame(self, frame): model = LocalsModel(self) model.update_frame(frame) self.setModel(model) -class Debugger(QObject,ComponentMixin): - name = 'Debugger' - - preferences = Parameter.create(name='Preferences',children=[ - {'name': 'Reload CQ', 'type': 'bool', 'value': False}, - {'name': 'Add script dir to path','type': 'bool', 'value': True}, - {'name': 'Change working dir to script dir','type': 'bool', 'value': True}, - {'name': 'Reload imported modules', 'type': 'bool', 'value': True}, - ]) +class Debugger(QObject, ComponentMixin): + name = "Debugger" + preferences = Parameter.create( + name="Preferences", + children=[ + {"name": "Reload CQ", "type": "bool", "value": False}, + {"name": "Add script dir to path", "type": "bool", "value": True}, + {"name": "Change working dir to script dir", "type": "bool", "value": True}, + {"name": "Reload imported modules", "type": "bool", "value": True}, + ], + ) sigRendered = pyqtSignal(dict) sigLocals = pyqtSignal(dict) - sigTraceback = pyqtSignal(object,str) + sigTraceback = pyqtSignal(object, str) sigFrameChanged = pyqtSignal(object) sigLineChanged = pyqtSignal(int) sigLocalsChanged = pyqtSignal(dict) - sigCQChanged = pyqtSignal(dict,bool) + sigCQChanged = pyqtSignal(dict, bool) sigDebugging = pyqtSignal(bool) - _frames : List[FrameType] + _frames: List[FrameType] - def __init__(self,parent): - - super(Debugger,self).__init__(parent) + def __init__(self, parent): + super(Debugger, self).__init__(parent) ComponentMixin.__init__(self) self.inner_event_loop = QEventLoop(self) - self._actions = \ - {'Run' : [QAction(icon('run'), - 'Render', - self, - shortcut='F5', - triggered=self.render), - QAction(icon('debug'), - 'Debug', - self, - checkable=True, - shortcut='ctrl+F5', - triggered=self.debug), - QAction(icon('arrow-step-over'), - 'Step', - self, - shortcut='ctrl+F10', - triggered=lambda: self.debug_cmd(DbgState.STEP)), - QAction(icon('arrow-step-in'), - 'Step in', - self, - shortcut='ctrl+F11', - triggered=lambda: self.debug_cmd(DbgState.STEP_IN)), - QAction(icon('arrow-continue'), - 'Continue', - self, - shortcut='ctrl+F12', - triggered=lambda: self.debug_cmd(DbgState.CONT)) - ]} - + self._actions = { + "Run": [ + QAction( + icon("run"), "Render", self, shortcut="F5", triggered=self.render + ), + QAction( + icon("debug"), + "Debug", + self, + checkable=True, + shortcut="ctrl+F5", + triggered=self.debug, + ), + QAction( + icon("arrow-step-over"), + "Step", + self, + shortcut="ctrl+F10", + triggered=lambda: self.debug_cmd(DbgState.STEP), + ), + QAction( + icon("arrow-step-in"), + "Step in", + self, + shortcut="ctrl+F11", + triggered=lambda: self.debug_cmd(DbgState.STEP_IN), + ), + QAction( + icon("arrow-continue"), + "Continue", + self, + shortcut="ctrl+F12", + triggered=lambda: self.debug_cmd(DbgState.CONT), + ), + ] + } + self._frames = [] def get_current_script(self): - - return self.parent().components['editor'].get_text_with_eol() + return self.parent().components["editor"].get_text_with_eol() def get_breakpoints(self): - - return self.parent().components['editor'].debugger.get_breakpoints() + return self.parent().components["editor"].debugger.get_breakpoints() def compile_code(self, cq_script): - try: - module = ModuleType('temp') - cq_code = compile(cq_script, '', 'exec') + module = ModuleType("temp") + cq_code = compile(cq_script, "", "exec") return cq_code, module except Exception: self.sigTraceback.emit(sys.exc_info(), cq_script) return None, None def _exec(self, code, locals_dict, globals_dict): - with ExitStack() as stack: - fname = self.parent().components['editor'].filename - p = Path(fname if fname else '').abspath().dirname() + fname = self.parent().components["editor"].filename + p = Path(fname if fname else "").abspath().dirname() - if self.preferences['Add script dir to path'] and p.exists(): - sys.path.insert(0,p) + if self.preferences["Add script dir to path"] and p.exists(): + sys.path.insert(0, p) stack.callback(sys.path.remove, p) - if self.preferences['Change working dir to script dir'] and p.exists(): + if self.preferences["Change working dir to script dir"] and p.exists(): stack.enter_context(p) - if self.preferences['Reload imported modules']: + if self.preferences["Reload imported modules"]: stack.enter_context(module_manager()) - exec(code, locals_dict, globals_dict) + exec(code, locals_dict, globals_dict) - def _rand_color(self, alpha = 0., cfloat=False): - #helper function to generate a random color dict - #for CQ-editor's show_object function + def _rand_color(self, alpha=0.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 + 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), - )} - - def _inject_locals(self,module): + (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), + ), + } + def _inject_locals(self, module): cq_objects = {} - def _show_object(obj,name=None, options={}): - + def _show_object( + obj, + name=None, + options={}, # all following inputs are ignored by cq-editor + parent=1, + clear=True, + port=3939, + axes=False, + axes0=False, + grid=False, + ticks=10, + ortho=True, + transparent=False, + default_color=(232, 176, 36), + reset_camera=True, + zoom=1.0, + default_edgecolor=(128, 128, 128), + render_edges=True, + render_normals=False, + render_mates=False, + mate_scale=1.0, + deviation=0.1, + angular_tolerance=0.2, + edge_accuracy=5.0, + ambient_intensity=1.0, + direct_intensity=0.12, + ): if name: - cq_objects.update({name : SimpleNamespace(shape=obj,options=options)}) + cq_objects.update({name: SimpleNamespace(shape=obj, options=options)}) else: - cq_objects.update({str(id(obj)) : SimpleNamespace(shape=obj,options=options)}) + cq_objects.update( + {str(id(obj)): SimpleNamespace(shape=obj, options=options)} + ) - def _debug(obj,name=None): + def _debug(obj, name=None): + _show_object(obj, name, options=dict(color="red", alpha=0.2)) - _show_object(obj,name,options=dict(color='red',alpha=0.2)) + module.__dict__["show_object"] = _show_object + module.__dict__["debug"] = _debug + module.__dict__["rand_color"] = self._rand_color + module.__dict__["log"] = lambda x: info(str(x)) + module.__dict__["cq"] = cq - module.__dict__['show_object'] = _show_object - module.__dict__['debug'] = _debug - module.__dict__['rand_color'] = self._rand_color - module.__dict__['log'] = lambda x: info(str(x)) - module.__dict__['cq'] = cq + return cq_objects, set(module.__dict__) - {"cq"} - return cq_objects, set(module.__dict__)-{'cq'} - - def _cleanup_locals(self,module,injected_names): - - for name in injected_names: module.__dict__.pop(name) + def _cleanup_locals(self, module, injected_names): + for name in injected_names: + module.__dict__.pop(name) @pyqtSlot(bool) def render(self): - seed(59798267586177) #reset the seed every time render is called (preserves colors run to run) - if self.preferences['Reload CQ']: + seed( + 59798267586177 + ) # reset the seed every time render is called (preserves colors run to run) + if self.preferences["Reload CQ"]: reload_cq() cq_script = self.get_current_script() - cq_code,module = self.compile_code(cq_script) + cq_code, module = self.compile_code(cq_script) - if cq_code is None: return + if cq_code is None: + return - cq_objects,injected_names = self._inject_locals(module) + cq_objects, injected_names = self._inject_locals(module) try: self._exec(cq_code, module.__dict__, module.__dict__) - #remove the special methods - self._cleanup_locals(module,injected_names) + # remove the special methods + self._cleanup_locals(module, injected_names) - #collect all CQ objects if no explicit show_object was called + # collect all CQ objects if no explicit show_object was called if len(cq_objects) == 0: cq_objects = find_cq_objects(module.__dict__) self.sigRendered.emit(cq_objects) - self.sigTraceback.emit(None, - cq_script) + self.sigTraceback.emit(None, cq_script) self.sigLocals.emit(module.__dict__) except Exception: exc_info = sys.exc_info() sys.last_traceback = exc_info[-1] self.sigTraceback.emit(exc_info, cq_script) - + @property def breakpoints(self): - return [ el[0] for el in self.get_breakpoints()] + return [el[0] for el in self.get_breakpoints()] @pyqtSlot(bool) - def debug(self,value): - + def debug(self, value): previous_trace = sys.gettrace() if value: @@ -284,79 +317,73 @@ class Debugger(QObject,ComponentMixin): self.state = DbgState.STEP self.script = self.get_current_script() - code,module = self.compile_code(self.script) + code, module = self.compile_code(self.script) if code is None: self.sigDebugging.emit(False) - self._actions['Run'][1].setChecked(False) + self._actions["Run"][1].setChecked(False) return - cq_objects,injected_names = self._inject_locals(module) + cq_objects, injected_names = self._inject_locals(module) - #clear possible traceback - self.sigTraceback.emit(None, - self.script) + # clear possible traceback + self.sigTraceback.emit(None, self.script) try: sys.settrace(self.trace_callback) - exec(code,module.__dict__,module.__dict__) + exec(code, module.__dict__, module.__dict__) except Exception: exc_info = sys.exc_info() sys.last_traceback = exc_info[-1] - self.sigTraceback.emit(exc_info, - self.script) + self.sigTraceback.emit(exc_info, self.script) finally: sys.settrace(previous_trace) self.sigDebugging.emit(False) - self._actions['Run'][1].setChecked(False) + self._actions["Run"][1].setChecked(False) if len(cq_objects) == 0: cq_objects = find_cq_objects(module.__dict__) self.sigRendered.emit(cq_objects) - self._cleanup_locals(module,injected_names) + self._cleanup_locals(module, injected_names) self.sigLocals.emit(module.__dict__) - + self._frames = [] else: sys.settrace(previous_trace) self.inner_event_loop.exit(0) - - def debug_cmd(self,state=DbgState.STEP): - + def debug_cmd(self, state=DbgState.STEP): self.state = state self.inner_event_loop.exit(0) - - def trace_callback(self,frame,event,arg): - + def trace_callback(self, frame, event, arg): filename = frame.f_code.co_filename - if filename==DUMMY_FILE: + if filename == DUMMY_FILE: if not self._frames: self._frames.append(frame) - self.trace_local(frame,event,arg) + self.trace_local(frame, event, arg) return self.trace_callback else: return None - def trace_local(self,frame,event,arg): - + def trace_local(self, frame, event, arg): lineno = frame.f_lineno if event in (DbgEevent.LINE,): - if (self.state in (DbgState.STEP, DbgState.STEP_IN) and frame is self._frames[-1]) \ - or (lineno in self.breakpoints): - + if ( + self.state in (DbgState.STEP, DbgState.STEP_IN) + and frame is self._frames[-1] + ) or (lineno in self.breakpoints): if lineno in self.breakpoints: self._frames.append(frame) - + self.sigLineChanged.emit(lineno) self.sigFrameChanged.emit(frame) self.sigLocalsChanged.emit(frame.f_locals) - self.sigCQChanged.emit(find_cq_objects(frame.f_locals),True) + self.sigCQChanged.emit(find_cq_objects(frame.f_locals), True) self.inner_event_loop.exec_() @@ -375,7 +402,7 @@ class Debugger(QObject,ComponentMixin): @contextmanager def module_manager(): - """ unloads any modules loaded while the context manager is active """ + """unloads any modules loaded while the context manager is active""" loaded_modules = set(sys.modules.keys()) try: diff --git a/cq_editor/widgets/object_tree.py b/cq_editor/widgets/object_tree.py index 3a610b5..004477f 100644 --- a/cq_editor/widgets/object_tree.py +++ b/cq_editor/widgets/object_tree.py @@ -1,391 +1,424 @@ -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, set_color -from .viewer import DEFAULT_FACE_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 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'), - ((0.2,0,0),'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 +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, + set_color, +) +from .viewer import DEFAULT_FACE_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.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": 0.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"), + ((0.2, 0, 0), "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={}, # all following inputs are ignored by cq-editor + parent=1, + clear=True, + port=3939, + axes=False, + axes0=False, + grid=False, + ticks=10, + ortho=True, + transparent=False, + default_color=(232, 176, 36), + reset_camera=True, + zoom=1.0, + default_edgecolor=(128, 128, 128), + render_edges=True, + render_normals=False, + render_mates=False, + mate_scale=1.0, + deviation=0.1, + angular_tolerance=0.2, + edge_accuracy=5.0, + ambient_intensity=1.0, + direct_intensity=0.12, + ): + 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