StoryCheckTools/graph/directed_acyclic_graph/DAGPresent.py

from graph.directed_acyclic_graph.DAGLayout import DAGGraph, DAGOrderHelper
from graph.DataType import Arrow, Point
from typing import List
from PyQt5.QtWidgets import QWidget, QApplication, QGraphicsView, QGraphicsScene, QGraphicsItem
from PyQt5.QtCore import QRectF, Qt, QPointF
from PyQt5.QtGui import QFont, QBrush, QPen, QPainter
import sys
from enum import Enum
from abc import abstractmethod


class GraphNode:
    @abstractmethod
    def boundingRect(self) -> QRectF:
        raise NotImplementedError("boundingRect")

    @abstractmethod
    def pos(self)->QPointF:
        raise NotImplementedError("pos")


class StoryFragmentNode(QGraphicsItem, GraphNode):
    """
    故事情节名称节点
    """
    def __init__(self, name:str, font: QFont, parent):
        QGraphicsItem.__init__(self, parent)

        self.fragment_name = name
        self.font_bind = font
        self.data_bind: object = None
        pass

    def boundingRect(self) -> QRectF:
        size = self.font_bind.pixelSize()
        return QRectF(0, 0, size * len(self.fragment_name) + 10, size + 10)

    def paint(self, painter, option, widget = ...):
        outline = self.boundingRect()
        text_rect = QRectF(outline.x()+5, outline.y()+5, outline.width()-10, outline.height() - 10)

        painter.save()
        painter.fillRect(outline, QBrush(Qt.gray))
        painter.drawRect(outline)
        painter.setFont(self.font_bind)
        painter.drawText(text_rect, self.fragment_name)
        painter.restore()

        pass


    # 添加自定义功能函数
    def attach_data(self, data_inst: object):
        self.data_bind = data_inst
        pass

    def get_data(self) -> object:
        return self.data_bind

    pass


class PenetrateNode(QGraphicsItem, GraphNode):
    """
    贯穿连线节点
    """
    def __init__(self, width:float, parent):
        QGraphicsItem.__init__(self, parent)
        self.width_store = width
        pass

    def resize(self, width: float):
        self.width_store = width
        self.update()
        pass

    def boundingRect(self):
        return QRectF(0, 0, self.width_store, 8)

    def paint(self, painter, option, widget = ...):
        outline = self.boundingRect()

        painter.save()
        painter.fillRect(QRectF(0, 2, outline.width(), 4), Qt.black)
        painter.restore()
        pass


class CurveTransition(QGraphicsItem):
    def __init__(self, start: GraphNode, end: GraphNode, parent):
        QGraphicsItem.__init__(self, parent)
        self.start_node = start
        self.end_node = end
        self.outline = QRectF()
        pass

    def layout_refresh(self):
        self.setPos(self.start_node.boundingRect().topRight())
        self.outline = QRectF(0, 0,
                              self.end_node.pos().x() - self.start_node.pos().x() - self.start_node.boundingRect().width(),
                              self.end_node.pos().y() + self.end_node.boundingRect().height() - self.start_node.pos().y())
        pass

    def boundingRect(self):
        return self.outline
        pass

    def paint(self, painter, option, widget = ...):
        outline = self.outline

        painter.save()
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        start_rect = self.start_node.boundingRect()
        end_rect = self.end_node.boundingRect()

        start_pos = QPointF(0, start_rect.height()/2)
        end_pos = QPointF(outline.width(), outline.height() - end_rect.height()/2)

        npen = QPen(Qt.black)
        npen.setWidthF(4)
        painter.setPen(npen)
        painter.drawLine(start_pos, end_pos)
        painter.restore()
        pass


class Direction(Enum):
    RankLR = 0,
    RankTB = 1,


class DAGActiveView(QGraphicsView):
    def __init__(self, orie: Direction, parent):
        QGraphicsView.__init__(self, parent)

        self.transition_with = 200

        self.orientation = orie
        self.scene_bind = QGraphicsScene(self)
        self.setScene(self.scene_bind)

        font_global = QFont()
        font_global.setPixelSize(20)

        one = StoryFragmentNode("实例节点：示例情节", font_global, None)
        self.scene_bind.addItem(one)
        two = PenetrateNode(200, None)
        two.setPos(QPointF(500, 400))
        self.scene_bind.addItem(two)

        connectx = CurveTransition(one, two, None)
        self.scene_bind.addItem(connectx)
        connectx.layout_refresh()

        pass

    def update_with_edges(self, arrows: List[Arrow]) -> None:
        tools = DAGGraph()
        tools.rebuild_from_edges(arrows)
        tools.graph_layout()

        total_nodes = tools.nodes_with_layout
        for layer_idx in range(0, tools.max_layer_count):
            current_level_nodes: List[DAGOrderHelper] = list(filter(lambda n:n.layer_number == layer_idx, total_nodes))
            current_level_nodes.sort(key=lambda n:n.sort_number)

            # 构建当前层节点
            ypos_acc = 0
            all_nodes = []
            for node in current_level_nodes:
                if node.is_fake_node():
                    node = PenetrateNode(20, None)
                    ypos_acc += node.boundingRect().height()
                    all_nodes.append(node)
                    self.scene_bind.addItem(node)
                else:
                    node = StoryFragmentNode(node.layer_bind.bind_point().point_name, self.font(), None)
                    ypos_acc += node.boundingRect().height()
                    all_nodes.append(node)
                    self.scene_bind.addItem(node)
                    pass
                pass

            width_max = 0
            for n in all_nodes:
                width_max = max(width_max, n.boundingRect().width())
                pass

            for n in all_nodes:
                if n is PenetrateNode:
                    n.resize(width_max)
                    pass
                pass
            pass


        pass


if __name__ == "__main__":
    app = QApplication(sys.argv)
    view = DAGActiveView(Direction.RankLR, None)
    view.show()
    app.exec()