modelverse/sketchUI/exm_mainwindow.py

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())