modelverse/sketchUI/exm_scene.py

from enum import Enum
from PyQt5.QtWidgets import QGraphicsScene, QGraphicsItem, QGraphicsLineItem, QGraphicsRectItem, \
    QGraphicsEllipseItem, QInputDialog, QGraphicsItemGroup
from PyQt5.Qt import Qt, QPointF, QPen, QTransform
from sketchUI.graphics_edge_item import GraphicsEdgeItem
from sketchUI.graphics_node_item import GraphicsNodeItem
from sketchUI import mvops
from evolution.node_ops import NodeAdd, NodeDelete, NodeRetype
from commons import all_nodes_with_type


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._connect_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):
        # type: (GraphicsNodeItem, GraphicsNodeItem, bool) -> None
        # draw an edge between two items. If is_new, also add it to the model
        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 or not isinstance(item, GraphicsNodeItem):
                return
            else:
                self._connect_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 or not isinstance(item, GraphicsNodeItem):
                    return
                self.draw_edge(self._connect_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:
            self._handle_keypress_delete(self.selectedItems())

        # "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)

        # "T" lets user type selected element
        elif event.key() == Qt.Key_T:
            # exactly one element that is a group must be selected
            selected = self.selectedItems()
            if len(selected) != 1:
                return
            item = selected[0]
            if isinstance(item, QGraphicsItemGroup):
                self._handle_keypress_type_on_group(item)
            elif isinstance(item, GraphicsNodeItem):
                self._handle_keypress_type_on_node(item)
            else:
                print("Cannot type element {}".format(item))

        else:
            QGraphicsScene.keyPressEvent(self, event)
        self.clearSelection()

    def _handle_keypress_type_on_node(self, item):
        # type: (GraphicsNodeItem) -> None
        # type an already typed node = retype it
        node_type, ok = QInputDialog.getText(self._parent, "Retype node", "New type", text=item.get_type())
        if not ok:
            # user canceled
            return
        if node_type:
            print("Reyping item {} to type {}".format(item, node_type))
            if node_type in mvops.get_available_types():
                print("Error: Already such a type: {}".format(node_type))
                return

        # local or global retype?
        scope, ok = QInputDialog.getItem(self._parent, "Select scope", "Scope", ["Local", "Global"])
        if not ok:
            return

        retype_handler = NodeRetype()
        if scope == "Global":
            retype_handler.execute(self._cur_model, item.node_id, node_type, local=False)
        else:
            retype_handler.execute(self._cur_model, item.node_id, node_type, local=True)

        # rename on screen
        for node_item in self.items():
            if not isinstance(node_item, GraphicsNodeItem):
                continue
            if item.get_type() == node_item.get_type():
                node_item.set_type(node_type)

        # update list widget
        self._parent.populate_types()

    def _handle_keypress_type_on_group(self, group):
        # type: (QGraphicsItemGroup) -> None
        # type the selected group = make a real node out of it and store it in the actual model
        # get the type from the user
        node_type, ok = QInputDialog.getText(self._parent, "Type node", "Enter type")
        if not ok:
            # user canceled
            return
        if node_type:
            print("Typing item {} to type {}".format(group, node_type))
            if node_type in mvops.get_available_types():
                print("Error: Already such a type: {}".format(node_type))
                return

        # perform add local or global?
        scope, ok = QInputDialog.getItem(self._parent, "Select scope", "Scope", ["Local", "Global"])
        if not ok:
            return

        # add the node to the model
        add_handler = NodeAdd()
        if scope == "Global":
            add_handler.execute(self._cur_model, node_type, local=False)
        else:
            add_handler.execute(self._cur_model, node_type, local=True)
        # Get node id of newly added node in current model
        nodeid = all_nodes_with_type(self._cur_model, node_type)[0]

        # update view
        nodeitem = GraphicsNodeItem(nodeid, node_type)
        nodeitem.setPos(group.x(), group.y())
        nodeitem.setFlag(QGraphicsItem.ItemIsSelectable, True)
        nodeitem.setFlag(QGraphicsItem.ItemIsMovable, True)
        self.removeItem(group)
        self.addItem(nodeitem)
        self._parent.populate_types()

    def _handle_keypress_delete(self, selected):
        del_hander = NodeDelete()
        for item in selected:
            # only delete nodes, edges are taken care of later
            if isinstance(item, GraphicsNodeItem):
                # when deleting a node, local or global?
                scope, ok = QInputDialog.getItem(self._parent, "Select scope", "Scope", ["Local", "Global"])
                if not ok:
                    return
                if scope == "Global":
                    # global language evolution, so delete node with same type everywhere
                    del_hander.execute(self._cur_model, item.node_id, local=False, check_if_last=False)
                else:
                    # just local, delete from this model only
                    del_hander.execute(self._cur_model, item.node_id, local=True, check_if_last=True)

                # in view, delete edges that were connected to this node as well
                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)
                # repopulate available types in view since they might have changed
            else:
                for item in selected:
                    self.removeItem(item)

        # if any node was deleted, repopulate list if available items
        if any(isinstance(item, GraphicsNodeItem) for item in selected):
            self._parent.populate_types()