modelverse/sketchUI/im_mainwindow.py

from PyQt5.QtWidgets import QMainWindow, QGraphicsItem, QAction, QActionGroup, QGraphicsView
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import QStateMachine, QState
from sketchUI.ui import Ui_MainWindow
from sketchUI.im_scene import CustomScene, Mode
from sketchUI import mvops
from sketchUI.graphics_node_item import GraphicsNodeItem
from wrappers.modelverse import element_list_nice


class IMMainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, model):
        QMainWindow.__init__(self)
        self.setupUi(self)
        self.setWindowTitle(model)
        self._cur_model = model

        self._scene = CustomScene(model)
        self._scene.set_mode(Mode.SELECT)
        self._scene.setSceneRect(0, 0, 200, 200)
        self.graphicsView.setScene(self._scene)

        self.listWidget.addItems(mvops.get_available_types())
        self.listWidget.itemDoubleClicked.connect(self._on_list_item_clicked)

        self.setup_toolbar()
        self.setup_state_machine()

        # load the model
        self._load_model()

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

    def setup_toolbar(self):
        self.select_action = QAction("Select", self)
        self.select_action.setIcon(QIcon("sketchUI/icons/select.png"))
        self.select_action.setCheckable(True)
        self.select_action.setChecked(True)
        self.connect_action = QAction("Connect", self)
        self.connect_action.setIcon(QIcon("sketchUI/icons/connect.png"))
        self.connect_action.setCheckable(True)
        action_group = QActionGroup(self)
        action_group.setExclusive(True)
        action_group.addAction(self.select_action)
        action_group.addAction(self.connect_action)

        for item in action_group.actions():
            self.toolBar.addAction(item)

    def setup_state_machine(self):
        self._statemachine = QStateMachine()
        state_select = QState()
        state_connect = QState()
        state_select.addTransition(self.connect_action.triggered, state_connect)
        state_connect.addTransition(self.select_action.triggered, state_select)
        state_connect.entered.connect(self._state_connect_entered)
        state_select.entered.connect(self._state_select_entered)

        self._statemachine.addState(state_select)
        self._statemachine.addState(state_connect)
        self._statemachine.setInitialState(state_select)

    def _load_model(self):
        """
        Load the model from the mv and render to screen
        Likely to be replaced later since rendering ideally already happens in the mv
        """
        model = element_list_nice(self._cur_model)
        if not model:
            # emtpy model
            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)

    def _state_connect_entered(self):
        self._scene.set_mode(Mode.CONNECT)
        self._make_items_movable(False)
        self._enable_list_widget(False)

    def _state_select_entered(self):
        self._scene.set_mode(Mode.SELECT)
        self._make_items_movable(True)
        self._enable_box_select(True)
        self._enable_list_widget(True)

    def _enable_list_widget(self, enabled):
        self.listWidget.setEnabled(enabled)

    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 _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 new node to model in mv
        node_id = mvops.add_node(self._cur_model, event.text())
        # render to scene
        self._add_node_to_scene(node_id, event.text())