from PyQt5.QtWidgets import QMainWindow, QAction, QActionGroup, QGraphicsItem, QGraphicsView from PyQt5.QtGui import QIcon from PyQt5.QtCore import QStateMachine, QState from sketchUI.ui import Ui_MainWindow from sketchUI.exm_scene import SketchScene, Mode from sketchUI import mvops from sketchUI.graphics_node_item import GraphicsNodeItem from wrappers.modelverse import element_list_nice class EXMMainWindow(QMainWindow, Ui_MainWindow): def __init__(self, model): QMainWindow.__init__(self) self.setupUi(self) self.setWindowTitle(model) self._cur_model = model self._scene = SketchScene(model, self) self._scene.setSceneRect(0, 0, 200, 200) self.graphicsView.setScene(self._scene) self._scene.set_mode(Mode.SELECT) self.populate_types() self.listWidget.itemDoubleClicked.connect(self._on_list_item_clicked) self._setup_toolbar() self._setup_statemachine() # load model self._load_model() # lastly, run state machine self._statemachine.start() def populate_types(self): # sync example model types to list widget self.listWidget.clear() self.listWidget.addItems(mvops.get_available_types()) def _setup_toolbar(self): # create action buttons on toolbar self.select_action = QAction("Select", self) self.select_action.setIcon(QIcon("sketchUI/icons/select.png")) self.select_action.setCheckable(True) self.connect_action = QAction("Connect", self) self.connect_action.setIcon(QIcon("sketchUI/icons/connect.png")) self.connect_action.setCheckable(True) self.line_action = QAction("Draw line", self) self.line_action.setIcon(QIcon("sketchUI/icons/line.png")) self.line_action.setCheckable(True) self.rect_action = QAction("Draw Rectangle", self) self.rect_action.setIcon(QIcon("sketchUI/icons/rect.png")) self.rect_action.setCheckable(True) self.circle_action = QAction("Draw Circle", self) self.circle_action.setIcon(QIcon("sketchUI/icons/circle.png")) self.circle_action.setCheckable(True) self.pen_action = QAction("Pen", self) self.pen_action.setIcon(QIcon("sketchUI/icons/pen.svg")) self.pen_action.setCheckable(True) action_group = QActionGroup(self) action_group.setExclusive(True) action_group.addAction(self.select_action) action_group.addAction(self.connect_action) action_group.addAction(self.line_action) action_group.addAction(self.rect_action) action_group.addAction(self.circle_action) action_group.addAction(self.pen_action) for item in action_group.actions(): self.toolBar.addAction(item) def _setup_statemachine(self): self._statemachine = QStateMachine() state_select = QState() state_connect = QState() state_draw_line = QState() state_draw_rect = QState() state_draw_circle = QState() state_draw_free = QState() # from select state_select.addTransition(self.select_action.triggered, state_select) state_select.addTransition(self.connect_action.triggered, state_connect) state_select.addTransition(self.line_action.triggered, state_draw_line) state_select.addTransition(self.rect_action.triggered, state_draw_rect) state_select.addTransition(self.circle_action.triggered, state_draw_circle) state_select.addTransition(self.pen_action.triggered, state_draw_free) state_select.entered.connect(self._state_select_entered) # from connect state_connect.addTransition(self.connect_action.triggered, state_connect) state_connect.addTransition(self.select_action.triggered, state_select) state_connect.addTransition(self.line_action.triggered, state_draw_line) state_connect.addTransition(self.rect_action.triggered, state_draw_rect) state_connect.addTransition(self.circle_action.triggered, state_draw_circle) state_connect.addTransition(self.pen_action.triggered, state_draw_free) state_connect.entered.connect(self._state_connect_entered) # from rectangle state_draw_rect.addTransition(self.rect_action.triggered, state_draw_rect) state_draw_rect.addTransition(self.connect_action.triggered, state_connect) state_draw_rect.addTransition(self.select_action.triggered, state_select) state_draw_rect.addTransition(self.line_action.triggered, state_draw_line) state_draw_rect.addTransition(self.circle_action.triggered, state_draw_circle) state_draw_rect.addTransition(self.pen_action.triggered, state_draw_free) state_draw_rect.entered.connect(self._state_draw_rect_entered) # from line state_draw_line.addTransition(self.line_action.triggered, state_draw_line) state_draw_line.addTransition(self.rect_action.triggered, state_draw_rect) state_draw_line.addTransition(self.connect_action.triggered, state_connect) state_draw_line.addTransition(self.select_action.triggered, state_select) state_draw_line.addTransition(self.circle_action.triggered, state_draw_circle) state_draw_line.addTransition(self.pen_action.triggered, state_draw_free) state_draw_line.entered.connect(self._state_draw_line_entered) # from circle state_draw_circle.addTransition(self.circle_action.triggered, state_draw_circle) state_draw_circle.addTransition(self.line_action.triggered, state_draw_line) state_draw_circle.addTransition(self.rect_action.triggered, state_draw_rect) state_draw_circle.addTransition(self.connect_action.triggered, state_connect) state_draw_circle.addTransition(self.select_action.triggered, state_select) state_draw_circle.addTransition(self.pen_action.triggered, state_draw_free) state_draw_circle.entered.connect(self._state_draw_circle_entered) # from freehand state_draw_free.addTransition(self.pen_action.triggered, state_draw_free) state_draw_free.addTransition(self.circle_action.triggered, state_draw_circle) state_draw_free.addTransition(self.line_action.triggered, state_draw_line) state_draw_free.addTransition(self.rect_action.triggered, state_draw_rect) state_draw_free.addTransition(self.connect_action.triggered, state_connect) state_draw_free.addTransition(self.select_action.triggered, state_select) state_draw_free.entered.connect(self._state_draw_free_entered) self._statemachine.addState(state_select) self._statemachine.addState(state_connect) self._statemachine.addState(state_draw_rect) self._statemachine.addState(state_draw_line) self._statemachine.addState(state_draw_circle) self._statemachine.addState(state_draw_free) self._statemachine.setInitialState(state_select) def _state_select_entered(self): self._scene.set_mode(Mode.SELECT) self._make_items_movable(True) self._enable_box_select(True) def _state_connect_entered(self): self._scene.set_mode(Mode.CONNECT) self._make_items_movable(False) self._enable_box_select(False) def _state_draw_rect_entered(self): self._scene.set_mode(Mode.RECT) self._make_items_movable(False) self._enable_box_select(False) def _state_draw_line_entered(self): self._scene.set_mode(Mode.LINE) self._make_items_movable(False) self._enable_box_select(False) def _state_draw_circle_entered(self): self._scene.set_mode(Mode.CIRCLE) self._make_items_movable(False) self._enable_box_select(False) def _state_draw_free_entered(self): self._scene.set_mode(Mode.FREE) self._make_items_movable(False) self._enable_box_select(False) def _make_items_movable(self, movable): for item in self._scene.items(): try: # hacky hack because of pythons isinstance fails due to import chaos # edges are not selectable or movable item.__hack__() item.setFlag(QGraphicsItem.ItemIsMovable, False) item.setFlag(QGraphicsItem.ItemIsSelectable, False) continue except AttributeError: pass item.setFlag(QGraphicsItem.ItemIsMovable, movable) item.setFlag(QGraphicsItem.ItemIsSelectable, movable) def _enable_box_select(self, enable): if enable: self.graphicsView.setDragMode(QGraphicsView.RubberBandDrag) else: self.graphicsView.setDragMode(QGraphicsView.NoDrag) def _load_model(self): model = element_list_nice(self._cur_model) if not model: # empty return for item in model: typ = item["type"] if typ == "Node": # first, draw all nodes node_type = item["typeID"] # workaround since get_attrs returns strings with " and list_nice doesnt if not node_type.startswith("\""): node_type = "\"" + node_type + "\"" self._add_node_to_scene(item["id"], node_type) for item in model: typ = item["type"] if typ == "Edge": target = item["__target"] src = item["__source"] self._add_edge_to_scene(src, target) def _add_node_to_scene(self, node_id, node_type, x=0, y=0): item = GraphicsNodeItem(node_id) item.setText(node_type) item.setPos(x, y) item.setFlag(QGraphicsItem.ItemIsMovable, True) item.setFlag(QGraphicsItem.ItemIsSelectable, True) self._scene.addItem(item) def _add_edge_to_scene(self, from_id, to_id): from_item = None to_item = None for item in self._scene.items(): try: node_id = item.node_id except AttributeError: # no node item, continue continue if node_id == from_id: from_item = item continue if node_id == to_id: to_item = item self._scene.draw_edge(from_item, to_item, is_new=False) def _on_list_item_clicked(self, event): # add node to model node_id = mvops.add_node(self._cur_model, event.text()) # render node self._add_node_to_scene(node_id, event.text())