modelverse/sketchUI/exm_scene.py

from enum import Enum
from PyQt5.QtWidgets import QGraphicsScene, QGraphicsItem, QGraphicsLineItem, QGraphicsRectItem, \
    QGraphicsEllipseItem, QInputDialog
from PyQt5.Qt import Qt, QPointF, QPen, QTransform
from PyQt5.QtCore import pyqtSignal
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
    LINE = 2
    RECT = 3
    CIRCLE = 4
    FREE = 5


class SketchScene(QGraphicsScene):
    def __init__(self, model, parent):
        QGraphicsScene.__init__(self)
        self._cur_drawing = False
        self._mode = None  # set from mainwindow on start
        self._from_item = None  # mouse pressed on this item if in connect mode
        self._cur_model = model

        self._orig_point = QPointF()
        self._free_draw_lines = []

        self._parent = parent

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

    def _in_draw_mode(self):
        # are we in one of the draw modes? (rect, line, free, ...)
        if self._mode == Mode.CONNECT or self._mode == Mode.SELECT:
            return False
        return True

    def draw_edge(self, from_item, to_item, is_new):
        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)

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton and self._in_draw_mode():
            # start drawing, save click point
            self._orig_point = event.scenePos()
            self._cur_drawing = True
        elif event.button() == Qt.LeftButton and self._mode == Mode.CONNECT:
            item = self.itemAt(event.scenePos(), QTransform())
            if not item:
                return
            self._from_item = item
        else:
            pass

        QGraphicsScene.mousePressEvent(self, event)

    def mouseMoveEvent(self, event):
        # if in freehand mode, draw lines from move movement
        if self._mode == Mode.FREE and self._cur_drawing:
            pt = event.scenePos()
            line = QGraphicsLineItem(self._orig_point.x(), self._orig_point.y(), pt.x(), pt.y())
            line.setPen(QPen(Qt.black, 2, Qt.SolidLine))
            self.addItem(line)

            self._orig_point = pt
            self._free_draw_lines.append(line)

        QGraphicsScene.mouseMoveEvent(self, event)

    def mouseReleaseEvent(self, event):
        if self._cur_drawing:
            end_point = event.scenePos()

            if self._mode == Mode.LINE:
                line = QGraphicsLineItem(self._orig_point.x(), self._orig_point.y(),
                                         end_point.x(), end_point.y())
                line.setPen(QPen(Qt.black, 2, Qt.SolidLine))
                self.addItem(line)

            elif self._mode == Mode.RECT:
                width = abs(end_point.x() - self._orig_point.x())
                height = abs(end_point.y() - self._orig_point.y())
                rect = QGraphicsRectItem(self._orig_point.x(), self._orig_point.y(),
                                         width, height)
                rect.setPen(QPen(Qt.black, 2, Qt.SolidLine))
                self.addItem(rect)

            elif self._mode == Mode.CIRCLE:
                width = abs(end_point.x() - self._orig_point.x())
                height = abs(end_point.y() - self._orig_point.y())
                ellipse = QGraphicsEllipseItem(self._orig_point.x(), self._orig_point.y(),
                                               width, height)
                ellipse.setPen(QPen(Qt.black, 2, Qt.SolidLine))
                self.addItem(ellipse)

            elif self._mode == Mode.FREE:
                line = QGraphicsLineItem(self._orig_point.x(), self._orig_point.y(),
                                         end_point.x(), end_point.y())
                line.setPen(QPen(Qt.black, 2, Qt.SolidLine))
                self.addItem(line)

                # group lines together
                self._free_draw_lines.append(line)
                group = self.createItemGroup(self._free_draw_lines)
                group.setFlag(QGraphicsItem.ItemIsSelectable, False)
                group.setFlag(QGraphicsItem.ItemIsMovable, False)
                del self._free_draw_lines[:]

            else:
                pass

            self._cur_drawing = False

        else:
            if 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

            elif self._mode == Mode.CONNECT:
                item = self.itemAt(event.scenePos(), QTransform())
                if not item:
                    return
                self.draw_edge(self._from_item, item, is_new=True)

            else:
                pass

        QGraphicsScene.mouseReleaseEvent(self, event)

    def keyPressEvent(self, event):
        if not self._mode == Mode.SELECT:
            return
        # "del" deletes all selected items
        if event.key() == Qt.Key_Delete:
            for item in self.selectedItems():
                self.removeItem(item)

        # "G" groups selected items
        elif event.key() == Qt.Key_G:
            selected = self.selectedItems()
            self.clearSelection()

            group = self.createItemGroup(selected)
            bb_rect = QGraphicsRectItem(group.boundingRect())
            bb_rect.setPen(QPen(Qt.gray, 1, Qt.DashLine))
            group.addToGroup(bb_rect)
            group.setFlag(QGraphicsItem.ItemIsSelectable, True)
            group.setFlag(QGraphicsItem.ItemIsMovable, True)

        elif event.key() == Qt.Key_T:
            selected = self.selectedItems()
            if len(selected) != 1:
                return
            item = selected[0]
            # "T" lets the user type the selected item
            text, ok = QInputDialog.getText(self._parent, "Text input", "Enter type")
            if ok and text:
                print("Typing item {} to type {}".format(item, text))
                if text in mvops.get_available_types():
                    print("Already such a type: {}".format(text))
                    return
                self._type_item(item, text)
        else:
            QGraphicsScene.keyPressEvent(self, event)

    def _type_item(self, item, typ):
        # typing means elevating it to a real node
        nid = mvops.add_node(self._cur_model, typ)
        nodeitem = GraphicsNodeItem(nid)
        nodeitem.setText(typ)
        nodeitem.setPos(item.x(), item.y())
        nodeitem.setFlag(QGraphicsItem.ItemIsSelectable, True)
        nodeitem.setFlag(QGraphicsItem.ItemIsMovable, True)
        self.removeItem(item)
        self.addItem(nodeitem)
        self._parent.populate_types()