modelverse/sketchUI/im_scene.py

from enum import Enum
from PyQt5.QtWidgets import QGraphicsScene, QGraphicsItem, QTableWidgetItem, QInputDialog
from PyQt5.Qt import Qt, QTransform, QApplication
from sketchUI.graphics_edge_item import GraphicsEdgeItem
from sketchUI.graphics_node_item import GraphicsNodeItem
from sketchUI import mvops
import commons

class Mode(Enum):
    SELECT = 0
    CONNECT = 1


class CustomScene(QGraphicsScene):
    def __init__(self, model, parent):
        QGraphicsScene.__init__(self)
        self._mode = None  # set by mainwindow at start
        self.connect_from_item = None  # store the item to draw the connecting line from
        self._cur_model = model  # the mv model we operate on
        self._parent = parent  # mainwindow parent

    def set_mode(self, mode):
        self._mode = mode

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton and self._mode == Mode.CONNECT:
            item = self.itemAt(event.scenePos(), QTransform())
            if not item:
                return
            # store the potential item to be connected from
            self.connect_from_item = item

        if event.button() == Qt.LeftButton and self._mode == Mode.SELECT:
            # load attributes for selected node
            item = self.itemAt(event.scenePos(), QTransform())
            if not item:
                return
            elif isinstance(item, GraphicsNodeItem):
                self._parent.plainTextEdit.appendPlainText("Selected node {}:{}".format(item.node_id, item.get_type()))
                self.highlight_node(item.node_id, Qt.blue)
                # load node attributes and display them in table
                self._parent.tableWidget.setRowCount(0)
                self._parent.tableWidget.blockSignals(True)
                attrs = commons.get_attributes_of_node(self._cur_model, item.node_id)
                for attr in attrs:
                    table_item_key = QTableWidgetItem(attr.key)
                    table_item_key.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
                    table_item_val = QTableWidgetItem(attr.val)
                    cur_row_cnt = self._parent.tableWidget.rowCount()
                    self._parent.tableWidget.insertRow(cur_row_cnt)
                    self._parent.tableWidget.setItem(cur_row_cnt, 0, table_item_key)
                    self._parent.tableWidget.setItem(cur_row_cnt, 1, table_item_val)
                self._parent.tableWidget.blockSignals(False)

            elif isinstance(item, GraphicsEdgeItem):
                self._parent.plainTextEdit.appendPlainText("Selected edge")
            else:
                pass

        QGraphicsScene.mousePressEvent(self, event)

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self._mode == Mode.CONNECT:
            item = self.itemAt(event.scenePos(), QTransform())
            if not item:
                return
            if item == self.connect_from_item:
                self._parent.plainTextEdit.appendPlainText("Error: Cannot connect same elements")
                return
            self.draw_edge(self.connect_from_item, item, is_new=True)

        if event.button() == Qt.LeftButton and self._mode == Mode.SELECT:
            item = self.itemAt(event.scenePos(), QTransform())
            if not item:
                return
            for item in self.items():
                if isinstance(item, GraphicsEdgeItem):
                    item.redraw()

        QGraphicsScene.mouseReleaseEvent(self, event)

    def keyPressEvent(self, event):
        if not self._mode == Mode.SELECT:
            return

        # del deletes elements
        if event.key() == Qt.Key_Delete:
            self._handle_keypress_delete(self.selectedItems())
        elif event.key() == Qt.Key_A:
            self._handle_keypress_attribute(self.selectedItems())
        else:
            pass

        QGraphicsScene.keyPressEvent(self, event)

    def draw_edge(self, from_item, to_item, is_new, edge_id=None):
        # type: (GraphicsNodeItem, GraphicsNodeItem, bool, str) -> None
        """
        Draw edge on screen between the node items "from_item" and "to_item".
        If is_new, additionally instantiate the edge in the model (if supported).
        """
        from_type = from_item.get_type()
        to_type = to_item.get_type()
        if is_new:
            QApplication.setOverrideCursor(Qt.WaitCursor)
            if not commons.is_edge_supported(from_type, to_type):
                self._parent.plainTextEdit.appendPlainText("Error: Edge not supported between types {} and {}".format(from_type, to_type))
                QApplication.restoreOverrideCursor()
                return

        line = GraphicsEdgeItem(from_item, to_item, edge_id)
        line.setFlag(QGraphicsItem.ItemIsMovable, False)
        line.setFlag(QGraphicsItem.ItemIsSelectable, False)
        self.addItem(line)
        line.redraw()

        if is_new:
            edge_id = mvops.add_edge(self._cur_model, from_item.node_id, to_item.node_id)
            line.edge_id = edge_id
            self._parent.plainTextEdit.appendPlainText("Added edge between {} and {} to model".format(from_type, to_type))

        QApplication.restoreOverrideCursor()


    def _handle_keypress_delete(self, selected):
        if not len(selected) == 1:
            return
        item = selected[0]

        if isinstance(item, GraphicsNodeItem):
            # delete node in model (also deletes edges connected to it in model)
            self._parent.plainTextEdit.appendPlainText("Deleting node of type {}".format(item.get_type()))
            mvops.delete_node(self._cur_model, item.node_id)

            # in view, delete edges that were connected to this node as well
            # modelverse does this on its own so do not delete edges explicitly here
            for edge in self.items():
                if not isinstance(edge, GraphicsEdgeItem):
                    continue
                if edge.from_item.node_id == item.node_id or edge.to_item.node_id == item.node_id:
                    self.removeItem(edge)

            self.removeItem(item)

        if isinstance(item, GraphicsEdgeItem):
            self._parent.plainTextEdit.appendPlainText("Deleting edge")
            mvops.delete_edge(self._cur_model, item.edge_id)
            self.removeItem(item)


    def _handle_keypress_attribute(self, selected):
        if not len(selected) == 1:
            return
        item = selected[0]
        if not isinstance(item, GraphicsNodeItem):
            return

        # quickly check if attributing of such a node is allowed
        QApplication.setOverrideCursor(Qt.WaitCursor)
        self._parent.plainTextEdit.appendPlainText("Checking if attributing nodes of type {} is allowed ...".format(item.get_type()))
        if not mvops.is_attributing_allowed(item.get_type()):
            self._parent.plainTextEdit.appendPlainText("Error: Not allowed".format(item.get_type()))
            QApplication.restoreOverrideCursor()
            return
        QApplication.restoreOverrideCursor()
        self._parent.plainTextEdit.appendPlainText("Yes")

        # ask user for key value
        key, ok = QInputDialog.getText(self._parent, "New attribute", "Key value")
        if not ok or not key:
            return

        # check if key value already used for this node
        attrs = commons.get_attributes_of_node(self._cur_model, item.node_id)
        for attr in attrs:
            if attr.key == key:
                self._parent.plainTextEdit.appendPlainText("Error: Already such a key for the node: {}".format(key))
                return

        self._parent.add_new_attribute(key)

    def highlight_node(self, node_id, color):
        for item in self.items():
            if not isinstance(item, GraphicsNodeItem):
                continue
            if item.node_id == node_id:
                item.set_highlighted(True, color)
            else:
                item.set_highlighted(False, color)
            item.update(item.boundingRect())