modelverse/sketchUI/graphics_node_item.py

"""
QGraphics representation class for a node element.
An object with a defined size (120x120) on screen which can either be
an image file or a set of primitives.
"""

from enum import Enum
import base64
from PyQt5.QtWidgets import QGraphicsItem, QStyleOptionGraphicsItem, QWidget
from PyQt5.Qt import QRectF, QPainter, QFont, Qt, QBrush, QColor, QImage


class IconType(Enum):
    IMAGE = 0
    PRIMITIVE = 1


class GraphicsNodeItem(QGraphicsItem):
    def __init__(self, node_id, node_type="", consyn=None):
        QGraphicsItem.__init__(self)
        self.node_id = node_id
        self._type = node_type
        self._consyn = consyn

        self._icon_type = None
        self._font = QFont()
        self._font.setPixelSize(12)

        self.setZValue(1.0)

    def set_type(self, node_type):
        self._type = node_type
        self.update(self.boundingRect())

    def update_consyn(self, consyn):
        self._consyn = consyn
        self.update(self.boundingRect())

    def get_type(self):
        return self._type

    def boundingRect(self):
        bbox = QRectF(0, 0, 120, 120)
        return bbox

    def paint(self, painter, option, widget=None):
        # type: (QPainter, QStyleOptionGraphicsItem, QWidget) -> None
        # draw type text and bounding rectangle
        painter.setFont(self._font)

        painter.setBrush(QBrush(Qt.white))
        painter.drawRect(self.boundingRect())

        text_box = QRectF(0, 100, 120, 20)
        painter.drawRect(text_box)
        painter.setBrush(QBrush(QColor(0, 0, 0, 0)))
        painter.drawText(text_box, Qt.AlignCenter | Qt.TextDontClip, self._type)

        # draw concrete syntax
        if not self._consyn:
            # nothing to draw, so leave empty
            return

        # first, determine if we have an actual image file or a primitive group
        for item in self._consyn:
            if item["type"] == "Icon":
                if item["is_primitive"]:
                    self._icon_type = IconType.PRIMITIVE
                else:
                    self._icon_type = IconType.IMAGE
                break

        if not self._icon_type:
            print("Could not determine CS icon type for type {}".format(self._type))
            return

        if self._icon_type == IconType.IMAGE:
            # an actual image file, so load and display it
            img_data_b64 = None
            for item in self._consyn:
                if item["type"] == "Image":
                    img_data_b64 = item["data"]
                    break
            if not img_data_b64:
                print("Could not find image data for type {}".format(self._type))
                return

            img_data = base64.b64decode(img_data_b64)
            qimg = QImage.fromData(img_data)
            painter.drawImage(QRectF(0, 0, 100, 100), qimg)

        else:
            # a set of primitives
            for item in self._consyn:
                item_type = item["type"]
                # ignore the "useless" Icon class
                if item_type == "Icon":
                    continue

                if item_type == "Rectangle":
                    rect = QRectF(item["x"], item["y"], item["width"], item["height"])
                    painter.drawRect(rect)
                elif item_type == "Ellipse":
                    painter.drawEllipse(item["x"], item["y"], item["width"], item["height"])
                elif item_type == "Line":
                    painter.drawLine(item["start_x"], item["start_y"], item["end_x"], item["end_y"])
                else:
                    print("No paint implemented for primitive {}".format(item_type))