modelverse/sketchUI/im_scene.py

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


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


class CustomScene(QGraphicsScene):
    def __init__(self, model):
        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

    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
            self.connect_from_item = item

        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:
                print("Not connecting 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 draw_edge(self, from_item, to_item, is_new):
        # type: (GraphicsNodeItem, GraphicsNodeItem, bool) -> None
        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):
                print("edge not supported between {} 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)