modelverse/sketchUI/exm_mainwindow.py

from PyQt5.QtWidgets import QMainWindow, QAction, QActionGroup, QGraphicsItem, QGraphicsView, QTableWidgetItem
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import QStateMachine, QState, Qt
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 sketchUI.graphics_edge_item import GraphicsEdgeItem
from wrappers.modelverse import element_list_nice
import commons


class EXMMainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, model):
        QMainWindow.__init__(self)
        self.setupUi(self)
        self.setWindowTitle(model)
        self.menuActionVerify.setDisabled(True)
        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()

        # setup log viewer
        self.plainTextEdit.setReadOnly(True)

        # setup table view for attributes
        self.tableWidget.setColumnCount(2)
        self.tableWidget.setHorizontalHeaderLabels(["Key", "Value"])
        self.tableWidget.horizontalHeader().setStretchLastSection(True)
        self.tableWidget.itemChanged.connect(self._on_attribute_edited)

        # lastly, run state machine
        self._statemachine.start()

    def populate_types(self):
        # sync example model types to list widget
        self.listWidget.clear()
        self.listWidget.addItems(commons.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():
            if isinstance(item, GraphicsEdgeItem):
                # edges are never movable, but selectable
                item.setFlag(QGraphicsItem.ItemIsMovable, False)
                item.setFlag(QGraphicsItem.ItemIsSelectable, True)
            else:
                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"]
                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, item["id"])

    def _add_node_to_scene(self, node_id, node_type, x=0, y=0):
        consyn = mvops.get_consyn_of(node_type)
        item = GraphicsNodeItem(node_id, node_type, consyn)
        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, edge_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, edge_id=edge_id)

    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())
        self.plainTextEdit.appendPlainText("Added node of type {} to model".format(event.text()))

    def _on_attribute_edited(self, item):
        # type: (QTableWidgetItem) -> None
        """ An attribute was edited, change it in the model but do not check (too expensive so
        checking is done by verify method on demand).
        If the new entered value is empty, delete the attribute.
        """
        row = self.tableWidget.row(item)
        attr_key = self.tableWidget.item(row, 0).text()
        attr_val = self.tableWidget.item(row, 1).text()
        node = self._scene.selectedItems()[0]

        if not attr_val:
            self.plainTextEdit.appendPlainText("Deleting attribute {}".format(attr_key))
            mvops.delete_attribute_from_node(self._cur_model, node.node_id, attr_key)
            self.tableWidget.removeRow(row)
        else:
            self.plainTextEdit.appendPlainText("Updating value of attribute {} to {}".format(attr_key, attr_val))
            mvops.update_attribute_val(self._cur_model, node.node_id, attr_key, attr_val)

    def add_new_attribute(self, key, val="unknown"):
        """
        Adds a new attribute to the view with key "key" and optional val. Also adds this attribute to the modelverse
        model.
        """
        selected_node = self._scene.selectedItems()[0]
        self.plainTextEdit.appendPlainText("Adding new attribute with key {} to node {}".format(key, selected_node.get_type()))
        self.tableWidget.blockSignals(True)
        table_item_key = QTableWidgetItem(key)
        table_item_key.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
        table_item_val = QTableWidgetItem(val)
        cur_row_cnt = self.tableWidget.rowCount()
        self.tableWidget.insertRow(cur_row_cnt)
        self.tableWidget.setItem(cur_row_cnt, 0, table_item_key)
        self.tableWidget.setItem(cur_row_cnt, 1, table_item_val)
        self.tableWidget.blockSignals(False)

        # add to modelverse
        mvops.add_attribute(self._cur_model, selected_node.node_id, key, val)