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 or not isinstance(item, GraphicsNodeItem):
                return

            self._parent.tableWidget.setRowCount(0)
            self._parent.tableWidget.blockSignals(True)
            self._parent.plainTextEdit.appendPlainText("Selected item of type {}".format(item.get_type()))
            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)

        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():
                try:
                    item.__hack__()  # hack to check if item is of edge type
                    item.redraw()
                except AttributeError:
                    continue

        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):
        # type: (GraphicsNodeItem, GraphicsNodeItem, bool) -> 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:
            if not mvops.is_edge_supported(from_type, to_type):
                self._parent.plainTextEdit.appendPlainText("Error: Edge not supported between types {} and {}".format(from_type, to_type))
                return

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

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

    def _handle_keypress_delete(self, selected):
        for item in selected:
            # delete node in model (also deletes edges in model)
            if isinstance(item, GraphicsNodeItem):
                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)

    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)