# -*- coding: utf-8 -*- import math from PyQt5.QtCore import ( Qt, QPoint, QPointF, QSizeF, QLineF, QRectF, pyqtSlot, pyqtSignal, ) from PyQt5.QtGui import ( QPainter, QColor, QBrush, QPen, QPainterPath, QPolygonF, ) from PyQt5.QtWidgets import ( QApplication,QGraphicsScene, QGraphicsView, QGraphicsItem, QGraphicsTextItem, ) from Model.Network.Node import Node from Model.Network.Edge import Edge from Model.Network.Graph import Graph class NodeItem(QGraphicsItem): Type = QGraphicsItem.UserType + 1 def __init__(self, node, graph_widget): super(NodeItem, self).__init__() self.node = node self.setPos(QPointF(self.node.pos.x, self.node.pos.y)) self.graph = graph_widget self.setFlag(QGraphicsItem.ItemIsMovable) self.setFlag(QGraphicsItem.ItemSendsGeometryChanges) self.setCacheMode(QGraphicsItem.DeviceCoordinateCache) self.setZValue(1) def type(self): return NodeItem.Type def boundingRect(self): adjust = 2.0 return QRectF(-10 - adjust, -10 - adjust, 20 + adjust, 20 + adjust) def shape(self): path = QPainterPath() path.addEllipse(-10, -10, 20, 20) return path def paint(self, painter, option, widget): painter.setPen(Qt.NoPen) # Select color in function of node position in graph and # status color = Qt.darkBlue if self.graph.selected_new_edge_src_node() == self: color = Qt.darkRed elif self.graph.selected_item() == self: color = Qt.red elif not self.graph.graph.is_enable_node(self.node): color = Qt.darkGray elif self.graph.graph.is_upstream_node(self.node): color = Qt.yellow elif self.graph.graph.is_downstream_node(self.node): color = Qt.green painter.setBrush(QBrush(color)) painter.drawEllipse(-10, -10, 20, 20) def itemChange(self, change, value): if change == QGraphicsItem.ItemPositionHasChanged: self.graph.node_change_position(value, self) self.graph.update_edges(self) return super(NodeItem, self).itemChange(change, value) def mousePressEvent(self, event): self.update() super(NodeItem, self).mousePressEvent(event) def mouseReleaseEvent(self, event): self.update() super(NodeItem, self).mouseReleaseEvent(event) class EdgeItem(QGraphicsItem): Type = QGraphicsItem.UserType + 2 def __init__(self, src_node_item, dest_node_item, edge, graph_widget): super(EdgeItem, self).__init__() self.src_node = src_node_item self.dest_node = dest_node_item self.edge = edge self.graph = graph_widget self.setAcceptedMouseButtons(Qt.NoButton) def type(self): return Edge.Type def boundingRect(self): # Rectangle of edge for display update pen_width = 2.0 extra = (pen_width + 5) / 2.0 return QRectF( self.src_node.pos(), QSizeF( self.dest_node.pos().x() - self.src_node.pos().x(), self.dest_node.pos().y() - self.src_node.pos().y() ) ).normalized().adjusted(-extra, -extra, extra, extra) def shape(self): # Shape around edge for mouse selection vec = self.dest_node.pos() - self.src_node.pos() vec = vec * (5 / math.sqrt(QPointF.dotProduct(vec, vec))) extra = QPointF(vec.y(), -vec.x()) result = QPainterPath(self.src_node.pos() - vec + extra) result.lineTo(self.src_node.pos() - vec - extra) result.lineTo(self.dest_node.pos() + vec - extra) result.lineTo(self.dest_node.pos() + vec + extra) result.closeSubpath() return result def paint(self, painter, option, widget): # Draw shape of the edge # color = QColor(Qt.white) # color.setAlpha(128) # painter.setBrush(color) # painter.drawPath(self.shape()) line = QLineF(self.src_node.pos(), self.dest_node.pos()) if line.length() == 0.0: return # Select color color = Qt.darkBlue if self.graph.selected_item() == self: color = Qt.red elif self.graph.current_edge() == self: color = Qt.blue elif not self.graph.graph.is_enable_edge(self.edge): color = Qt.darkGray painter.setPen(QPen(color, 2, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)) # Draw the line painter.drawLine(line) line_center = QPointF( (self.src_node.pos().x() + self.dest_node.pos().x()) / 2, (self.src_node.pos().y() + self.dest_node.pos().y()) / 2, ) # Draw the arrows brush = QBrush() brush.setColor(color) brush.setStyle(Qt.SolidPattern) angle = math.acos(line.dx() / line.length()) if line.dy() >= 0: angle = (math.pi * 2.0) - angle size = 10.0 arrow_p1 = line_center + QPointF( math.sin(angle - math.pi / 3) * size, math.cos(angle - math.pi / 3) * size ) arrow_p2 = line_center + QPointF( math.sin(angle - math.pi + math.pi / 3) * size, math.cos(angle - math.pi + math.pi / 3) * size ) poly = QPolygonF([line_center, arrow_p1, arrow_p2]) path = QPainterPath() path.addPolygon(poly) painter.drawPolygon(poly) painter.fillPath(path, brush) class NodeText(QGraphicsTextItem): def __init__(self, node_item): super(NodeText, self).__init__() self.item = node_item self.setPlainText(self.item.node.name) self.setDefaultTextColor(Qt.black) self.set_custom_pos(self.item.pos()) self.setZValue(2) def set_custom_pos(self, pos): x = pos.x() y = pos.y() - 42 # Dont panic! The answer is 42 self.setPos(QPointF(x,y)) def paint(self, painter, option, widget): color = QColor(Qt.white) color.setAlpha(128) # Draw text background painter.setBrush(color) painter.drawRect(self.boundingRect()) # Draw text super(NodeText, self).paint(painter, option, widget) def rename(self): # Update the node text self.setPlainText(self.item.node.name) class NewEdgeLine(QGraphicsItem): def __init__(self, src, dest): super(NewEdgeLine, self).__init__() self.src = src self.dest = dest def boundingRect(self): pen_width = 2.0 extra = (pen_width + 5) / 2.0 return QRectF( self.src, QSizeF( self.dest.x() - self.src.x(), self.dest.y() - self.src.y() ) ).normalized().adjusted(-extra, -extra, extra, extra) def paint(self, painter, option, widget): line = QLineF(self.src, self.dest) if line.length() == 0.0: return color = Qt.darkGray painter.setPen(QPen(color, 2, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)) painter.drawLine(line) class GraphWidget(QGraphicsView): changeEdge = pyqtSignal(object) changeNode = pyqtSignal(object) def __init__(self, graph, parent=None): super(GraphWidget, self).__init__(parent=parent) self.timerId = 0 self.parent = parent self._state = "move" self._selected_item = None self._selected_new_edge_src_node = None self._current_edge = None self.tmp_line = None self.node_items = [] self.edge_items = [] self.texts = {} self.m_origin_x = 0.0 self.m_origin_y = 0.0 self.clicked = False scene = QGraphicsScene(self) scene.setItemIndexMethod(QGraphicsScene.NoIndex) scene.setSceneRect(0, 0, 2000, 2000) self.setScene(scene) self.setCacheMode(QGraphicsView.CacheBackground) self.setViewportUpdateMode(QGraphicsView.BoundingRectViewportUpdate) self.setRenderHint(QPainter.Antialiasing) self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse) self.setResizeAnchor(QGraphicsView.AnchorViewCenter) self.graph = graph self.scale(1, 1) self.previousScale = 1 self.setMinimumSize(400, 400) self.create_items() def create_items(self): """Create all items and draw its Returns: Nothing """ self.node_items = [] self.edge_items = [] self.texts = {} for node in self.graph.nodes(): inode = NodeItem(node, self) self.texts[inode] = NodeText(inode) self.scene().addItem(self.texts[inode]) self.scene().addItem(inode) self.node_items.append(inode) for edge in self.graph.edges(): n1 = list( filter( lambda n: n.node.name == edge.node1.name, self.node_items ) ) n2 = list( filter( lambda n: n.node.name == edge.node2.name, self.node_items ) ) iedge = EdgeItem(n1[0], n2[0], edge, self) self.scene().addItem(iedge) self.edge_items.append(iedge) def state(self, status): """Set the current status of draw widget Set the current status of draw widget, like "move", "add" or "del" Args: status: String of current status Returns: Nothing """ self._state = status def add_node(self, pos): """Create a new node in graph and draw it Args: pos: The position of new node Returns: Nothing """ node = self.graph.add_node(pos.x(), pos.y()) inode = NodeItem(node, self) self.scene().addItem(inode) self.node_items.append(inode) self.texts[inode] = NodeText(inode) self.scene().addItem(self.texts[inode]) self.changeNode.emit(self.sender()) def del_node(self, item): """Delete a node and update display Args: item: The node item to delete Returns: Nothing """ node = item.node edges = list( filter( lambda ie: (ie.edge.node1 == node or ie.edge.node2 == node), self.edge_items ) ) self.set_selected_item(None) for edge in edges: self.graph.remove_edge(edge.edge.name) self.graph.remove_node(node.name) self.scene().clear() self.create_items() self.changeNode.emit(self.sender()) self.changeEdge.emit(self.sender()) def rename_nodes(self): """Update all nodes item name in scene Returns: Nothing """ for i in self.texts: if type(i) == NodeItem: self.texts[i].rename() def display_update(self): """Clear the scene and redraw it Returns: Nothing """ self.scene().clear() self.create_items() def add_edge(self, node1, node2): """Create a new edge in graph and draw it Args: node1: The source node node2: The destination node Returns: Nothing """ if node1 == node2: return edge = self.graph.add_edge(node1.node, node2.node) self.set_selected_item(None) self.set_selected_new_edge_src_node(None) # Reset the temporary line self.tmp_line = None # Clear all scene and redraw it self.scene().clear() self.create_items() self.changeEdge.emit(self.sender()) def del_edge(self, item): """Delete an edge and update display Args: item: The edge item to delete Returns: Nothing """ edge = item.edge self.graph.remove_edge(edge.name) # Clear all scene and redraw it self.scene().clear() self.create_items() self.changeNode.emit(self.sender()) self.changeEdge.emit(self.sender()) def update_edges(self, node): """Update display of all edges linked with the node Args: node: The node item linked with edges to update Returns: Nothing """ edges = list( filter( lambda ie: (ie.edge.node1 == node.node or ie.edge.node2 == node.node), self.edge_items ) ) for edge in edges: edge.update() def node_change_position(self, pos, node): """Update node position and node text position Args: pos: The new position node: The node item Returns: Nothing """ node.node.setPos(pos.x(), pos.y()) self.texts[node].set_custom_pos(pos) self.texts[node].update() def selected_item(self): """Current selected item Returns: Item if item ar selected, otherelse None """ return self._selected_item def set_selected_item(self, item): """Set current selected item Args: item: The new item to select Returns: Nothing """ previous_item = self._selected_item self._selected_item = item if previous_item: previous_item.update() if item: item.update() def selected_new_edge_src_node(self): """The current node item selected to add new edge Returns: Item if item ar selected, otherelse None """ return self._selected_new_edge_src_node def set_selected_new_edge_src_node(self, node): """Set the current node item selected to add new edge Args: node: The new node to select Returns: Nothing """ previous_node = self._selected_new_edge_src_node self._selected_new_edge_src_node = node if previous_node: previous_node.update() def current_edge(self): """The current selected edge Returns: Item if edge are selected, otherelse None """ return self._current_edge def set_current_edge(self, edge): """Set the current edge item selected Args: edge: The new edge to select Returns: Nothing """ previous_edge = self._current_edge self._current_edge = edge self.graph.current_reach(edge.edge) if previous_edge: previous_edge.update() def reset_selection(self): """Reset all the selected items Returns: Nothing """ self.set_selected_new_edge_src_node(None) if self.tmp_line is not None: self.tmp_line = None self.scene().clear() self.create_items() def keyPressEvent(self, event): key = event.key() if key == Qt.Key_Plus: self.scaleView(1.2) elif key == Qt.Key_Minus: self.scaleView(1 / 1.2) elif key == Qt.Key_Escape: self.reset_selection() else: super(GraphWidget, self).keyPressEvent(event) def drawBackground(self, painter, rect): sceneRect = self.sceneRect() painter.fillRect( rect.intersected(sceneRect), QBrush(Qt.lightGray) ) painter.setBrush(Qt.NoBrush) painter.drawRect(sceneRect) def wheelEvent(self, event): self.scaleView(math.pow(2.0, -event.angleDelta().y() / 240.0)) def scaleView(self, scaleFactor): factor = self.transform().scale( scaleFactor, scaleFactor ).mapRect(QRectF(0, 0, 1, 1)).width() if factor < 0.05 or factor > 1.5: return self.scale(scaleFactor, scaleFactor) def mousePressEvent(self, event): pos = self.mapToScene(event.pos()) self.clicked = True # Move item and select edge item if self._state == "move": self._selected_new_edge_src_node = None self.mouse_origin_x = pos.x() self.mouse_origin_y = pos.y() items = self.items(event.pos()) if items and type(items[0]) == EdgeItem: edge = items[0] if edge: self.set_current_edge(edge) # Add nodes and edges elif self._state == "add": items = self.items(event.pos()) nodes = list(filter(lambda i: type(i) == NodeItem, items)) if not nodes: self.add_node(pos) else: if self.selected_new_edge_src_node() is None: self.set_selected_new_edge_src_node(nodes[0]) else: self.add_edge(self.selected_new_edge_src_node(), nodes[0]) # Delete nodes and edges elif self._state == "del": self._selected_new_edge_src_node = None items = list( filter( lambda i: type(i) == NodeItem or type(i) == EdgeItem, self.items(event.pos()) ) ) if len(items) > 0: item = items[0] if type(item) == NodeItem: self.del_node(item) elif type(item) == EdgeItem: self.del_edge(item) self.update() super(GraphWidget, self).mousePressEvent(event) def mouseReleaseEvent(self, event): self.clicked = False self.update() super(GraphWidget, self).mouseReleaseEvent(event) def mouseMoveEvent(self, event): pos = self.mapToScene(event.pos()) # Selecte item on the fly items = self.items(event.pos()) selectable_items = list( filter( lambda i: (type(i) == NodeItem or type(i) == EdgeItem), items ) ) if selectable_items: self.set_selected_item(selectable_items[0]) elif self.clicked == False: self.set_selected_item(None) # Update temporary line if self.selected_new_edge_src_node() is not None: if self.tmp_line is None: self.tmp_line = NewEdgeLine( self.selected_new_edge_src_node().pos(), pos ) self.scene().addItem(self.tmp_line) else: self.tmp_line.dest = pos self.tmp_line.update() # If state is "move" if self._state == "move": # Move on scene if (not self.selected_item() and event.buttons() & Qt.LeftButton): old_p = self.mapToScene(self.m_origin_x, self.m_origin_y) new_p = self.mapToScene(event.pos()) translation = new_p - old_p self.translate(translation.x(), translation.y()) self.m_origin_x = event.x() self.m_origin_y = event.y() # Propagate event self.update() super(GraphWidget, self).mouseMoveEvent(event)