Node

  1# Node.py
  2
  3from PySide6.QtWidgets import QGraphicsItem, QGraphicsEllipseItem
  4from PySide6.QtCore import QRectF, Qt, QPointF
  5from PySide6.QtGui import QPainter, QPen, QBrush
  6from Edge import Edge
  7from NodeData import NodeData
  8from NodeLayout import NodeLayout
  9
 10class Node(QGraphicsItem):
 11    def __init__(self, node_data: NodeData):
 12        super().__init__()
 13        self.data = node_data
 14        self.setFlag(QGraphicsItem.ItemIsMovable)
 15        self.setFlag(QGraphicsItem.ItemIsSelectable)
 16        self.setFlag(QGraphicsItem.ItemSendsGeometryChanges)
 17        self.rect = QRectF(0, 0, self.data.width, self.data.height)
 18        self.resizeHandleSize = 10
 19        self.resizing = False
 20        self.resizeDirection = None
 21        self.content = NodeLayout(self)
 22        self.input_port = Port(self, QPointF(0, 25), "input")
 23        self.output_port = Port(self, QPointF(self.rect.width(), 25), "output")
 24        self.setPos(self.data.pos_x, self.data.pos_y)  # Set initial position from NodeData
 25        self.setAcceptHoverEvents(True)
 26        self.hovered = False  # Track hover state
 27
 28    def setWidth(self, width):
 29        self.rect.setWidth(width)
 30        self.output_port.setPos(width, 25)
 31        self.prepareGeometryChange()
 32        self.content.update_proxy_widget_geometry()  # Ensure the content size is updated
 33        self.data.width = width
 34
 35    def setHeight(self, height):
 36        self.rect.setHeight(height)
 37        self.prepareGeometryChange()
 38        self.content.update_proxy_widget_geometry()  # Ensure the content size is updated
 39        self.data.height = height
 40
 41    def boundingRect(self):
 42        return self.rect
 43
 44    def paint(self, painter, option, widget):
 45        pen = QPen(Qt.black, 2)
 46        if self.hovered:
 47            pen.setColor(Qt.red)
 48        painter.setPen(pen)
 49        painter.drawRect(self.rect)
 50        self.content.paint(painter, option, widget)  # Delegate the painting to NodeLayout
 51
 52    def itemChange(self, change, value):
 53        if change == QGraphicsItem.ItemPositionChange:
 54            for port in [self.input_port, self.output_port]:
 55                for edge in port.edges:
 56                    edge.update_position()
 57            # Sync position with NodeData
 58            self.data.pos_x = value.x()
 59            self.data.pos_y = value.y()
 60        return super().itemChange(change, value)
 61
 62    def mousePressEvent(self, event):
 63        pos = event.pos()
 64        if pos.x() >= self.rect.right() - self.resizeHandleSize and \
 65                pos.y() >= self.rect.bottom() - self.resizeHandleSize:
 66            self.resizing = True
 67            self.resizeDirection = "bottom_right"
 68            self.setCursor(Qt.SizeFDiagCursor)
 69        else:
 70            self.resizing = False
 71            super().mousePressEvent(event)
 72
 73    def mouseMoveEvent(self, event):
 74        if self.resizing:
 75            delta = event.pos() - event.lastPos()
 76            if self.resizeDirection == "bottom_right":
 77                self.setWidth(self.rect.width() + delta.x())
 78                self.setHeight(self.rect.height() + delta.y())
 79        else:
 80            super().mouseMoveEvent(event)
 81
 82    def mouseReleaseEvent(self, event):
 83        if self.resizing:
 84            self.resizing = False
 85            self.setCursor(Qt.ArrowCursor)
 86        else:
 87            super().mouseReleaseEvent(event)
 88
 89    def remove_node(self):
 90        # Remove all edges connected to this node
 91        for edge in self.input_port.edges[:]:
 92            edge.remove()
 93        for edge in self.output_port.edges[:]:
 94            edge.remove()
 95
 96        # Update prevs and nexts of connected nodes
 97        for prev_id in self.data.prevs:
 98            prev_node = self.scene().get_node_by_id(prev_id)
 99            if prev_node:
100                prev_node.data.nexts.remove(self.data.uniq_id)
101
102        for next_id in self.data.nexts:
103            next_node = self.scene().get_node_by_id(next_id)
104            if next_node:
105                next_node.data.prevs.remove(self.data.uniq_id)
106
107        # Finally, remove this node from the scene
108        self.scene().removeItem(self)
109
110    def hoverEnterEvent(self, event):
111        self.hovered = True
112        self.update()  # Trigger a repaint
113        super().hoverEnterEvent(event)
114
115    def hoverLeaveEvent(self, event):
116        self.hovered = False
117        self.update()  # Trigger a repaint
118        super().hoverLeaveEvent(event)
119
120class Port(QGraphicsEllipseItem):
121    def __init__(self, parent, position, port_type):
122        super().__init__(-5, -5, 10, 10, parent)
123        self.setBrush(Qt.black)
124        self.setFlag(QGraphicsItem.ItemIsMovable, False)
125        self.setPos(position)
126        self.port_type = port_type
127        self.edges = []
128
129    def mousePressEvent(self, event):
130        if self.port_type == "output":
131            edge = Edge(self)
132            self.edges.append(edge)
133            self.scene().addItem(edge)
134
135    def mouseMoveEvent(self, event):
136        if self.edges:
137            self.edges[-1].update_position(event.scenePos())
138
139    def mouseReleaseEvent(self, event):
140        if self.edges:
141            items = self.scene().items(event.scenePos())
142            for item in items:
143                if isinstance(item, Port) and item != self:
144                    if self.port_type == "output" and item.port_type == "input":
145                        self.edges[-1].set_destination(item)
146                        item.edges.append(self.edges[-1])
147                        break
148            else:
149                self.scene().removeItem(self.edges[-1])
150                self.edges.pop()