# -*- coding: utf-8 -*-
import math
from model.network.Node import Node
from model.network.Edge import Edge
from model.network.Graph import Graph
from PyQt5.QtCore import (
qAbs, QLineF, QPointF, QPoint, qrand, QRectF, QSizeF, qsrand,
Qt, QTime, pyqtSlot, pyqtSignal,
)
from PyQt5.QtGui import (
QBrush, QColor, QLinearGradient, QPainter,
QPainterPath, QPen, QPolygonF, QRadialGradient,
QFont,
)
from PyQt5.QtWidgets import (
QApplication, QGraphicsItem, QGraphicsScene,
QGraphicsView, QStyle, QGraphicsTextItem,
)
class LandMark(QGraphicsItem):
def __init__(self, width, height, parent=None):
super(LandMark, self).__init__()
self.width = width
self.height = height
self.center = QPoint(width / 2, height / 2)
self.parent = parent
def paint(self, painter, option, widget):
painter.setPen(Qt.NoPen)
painter.setPen(QPen(Qt.black, 0))
painter.drawLine(
-10 + self.center.x(), 0 + self.center.y(),
10 + self.center.x(), 0 + self.center.y()
)
painter.drawLine(
0 + self.center.x(), -10 + self.center.y(),
0 + self.center.x(), 10 + self.center.y()
)
def boundingRect(self):
return QRectF(
-10 + self.center.x(), -10 + self.center.y(),
10 + self.center.x(), 10 + self.center.y()
)
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.newPos = QPointF()
self.setFlag(QGraphicsItem.ItemIsMovable)
self.setFlag(QGraphicsItem.ItemSendsGeometryChanges)
self.setCacheMode(QGraphicsItem.DeviceCoordinateCache)
self.setZValue(1)
def type(self):
return NodeItem.Type
def advance(self):
if self.newPos == self.pos():
return False
self.setPos(self.newPos)
return True
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)
color = Qt.yellow
if self.graph.selectedNewEdgeSrcNode() == self:
color = Qt.darkRed
elif self.graph.selectedItem() == self:
color = Qt.red
elif self.graph.graph.is_upstream_node(self.node):
color = Qt.blue
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.itemMoved()
self.graph.nodeChangePosition(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.src_pos = src_node_item.pos()
self.dest_pos = dest_node_item.pos()
self.setAcceptedMouseButtons(Qt.NoButton)
self.adjust()
def type(self):
return Edge.Type
def sourceNode(self):
return self.src_node.node
def setSourceNode(self, node_item):
self.src_node = node_item
self.adjust()
def destNode(self):
return self.dest
def setDestNode(self, node_item):
self.dest_node = node_item
self.adjust()
def adjust(self):
line = QLineF(
self.mapFromItem(self.src_node, 0, 0),
self.mapFromItem(self.dest_node, 0, 0)
)
length = line.length()
self.prepareGeometryChange()
def boundingRect(self):
pen_width = 3.0
extra = (pen_width + 5) / 2.0
return QRectF(
self.src_pos,
QSizeF(
self.dest_pos.x() - self.src_pos.x(),
self.dest_pos.y() - self.src_pos.y()
)
).normalized().adjusted(-extra, -extra, extra, extra)
def paint(self, painter, option, widget):
# Draw the line itself.
line = QLineF(self.src_node.pos(), self.dest_node.pos())
if line.length() == 0.0:
return
color = Qt.blue
if self.graph.selectedItem() == self:
color = Qt.red
painter.setPen(QPen(color, 3, Qt.SolidLine, Qt.RoundCap,
Qt.RoundJoin))
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 if there's enough room.
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
)
painter.drawPolygon(QPolygonF([line_center, arrow_p1, arrow_p2]))
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.setPos(self.item.pos())
self.set_custom_pos(self.item.pos())
self.setZValue(2)
def set_custom_pos(self, pos):
x = pos.x()
y = pos.y() - 42
self.setPos(QPointF(x,y))
def paint(self, painter, option, widget):
color = QColor(Qt.white)
color.setAlpha(128)
painter.setBrush(color)
painter.drawRect(self.boundingRect())
super(NodeText, self).paint(painter, option, widget)
def rename(self):
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._selectedItem = None
self._selectedNewEdgeSrcNode = None
self.tmp_line = None
self.node_items = []
self.edge_items = []
self.texts = {}
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):
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)
# self.mark = LandMark(1000, 1000, self)
# self.scene().addItem(self.mark)
def state(self, status):
self._state = status
def add_node(self, pos):
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):
node = item.node
edges = list(
filter(
lambda ie: ie.edge.node1 == node or ie.edge.node2 == node,
self.edge_items
)
)
self.setSelectedItem(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):
for i in self.texts:
if type(i) == NodeItem:
self.texts[i].rename()
def add_edge(self, node1, node2):
if node1 == node2:
return
edge = self.graph.add_edge(node1.node, node2.node)
# iedge = EdgeItem(node1, node2, edge, self)
# self.scene().addItem(iedge)
# self.edge_items.append(iedge)
self.setSelectedItem(None)
self.setSelectedNewEdgeSrcNode(None)
self.tmp_line = None
self.scene().clear()
self.create_items()
self.changeEdge.emit(self.sender())
def del_edge(self, item):
edge = item.edge
self.graph.remove_edge(edge.name)
self.scene().clear()
self.create_items()
self.changeNode.emit(self.sender())
self.changeEdge.emit(self.sender())
def update_edges(self, node):
edges = list(
filter(
lambda ie: ie.edge.node1 == node.node or ie.edge.node2 == node.node,
self.edge_items
)
)
for edge in edges:
# edge.node1.update()
# edge.node2.update()
edge.update()
def itemMoved(self):
if not self.timerId:
self.timerId = self.startTimer(1000 / 25)
def nodeChangePosition(self, pos, node):
node.node.setPos(pos.x(), pos.y())
self.texts[node].set_custom_pos(pos)
self.texts[node].update()
def selectedItem(self):
return self._selectedItem
def setSelectedItem(self, item):
previous_item = self._selectedItem
self._selectedItem = item
if previous_item:
previous_item.update()
if item:
item.update()
def selectedNewEdgeSrcNode(self):
return self._selectedNewEdgeSrcNode
def setSelectedNewEdgeSrcNode(self, node):
previous_node = self._selectedNewEdgeSrcNode
self._selectedNewEdgeSrcNode = node
if previous_node:
previous_node.update()
def reset_selection(self):
self.setSelectedNewEdgeSrcNode(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()
elif key == Qt.Key_Space or key == Qt.Key_Enter:
for item in self.scene().items():
if isinstance(item, Node):
item.setPos(-150 + qrand() % 300, -150 + qrand() % 300)
else:
super(GraphWidget, self).keyPressEvent(event)
def drawBackground(self, painter, rect):
sceneRect = self.sceneRect()
gradient = QLinearGradient(sceneRect.topLeft(), sceneRect.bottomRight())
gradient.setColorAt(0, Qt.lightGray)
gradient.setColorAt(1, Qt.lightGray)
painter.fillRect(rect.intersected(sceneRect), QBrush(gradient))
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
# print(f"{scaleFactor} : {factor}")
self.scale(scaleFactor, scaleFactor)
def mousePressEvent(self, event):
pos = self.mapToScene(event.pos())
if self._state == "move":
self.mouse_origin_x = pos.x()
self.mouse_origin_y = pos.y()
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.selectedNewEdgeSrcNode() is None:
self.setSelectedNewEdgeSrcNode(nodes[0])
else:
self.add_edge(self.selectedNewEdgeSrcNode(), nodes[0])
elif self._state == "del":
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.update()
super(GraphWidget, self).mouseReleaseEvent(event)
def mouseMoveEvent(self, event):
pos = self.mapToScene(event.pos())
items = self.items(event.pos())
selectable_items = list(
filter(
lambda i: type(i) == NodeItem or type(i) == EdgeItem,
items
)
)
if selectable_items:
self.setSelectedItem(selectable_items[0])
else:
self.setSelectedItem(None)
if self.selectedNewEdgeSrcNode() is not None:
if self.tmp_line is None:
self.tmp_line = NewEdgeLine(
self.selectedNewEdgeSrcNode().pos(),
pos
)
self.scene().addItem(self.tmp_line)
else:
self.tmp_line.dest = pos
self.tmp_line.update()
if self._state == "move":
self.update()
super(GraphWidget, self).mouseMoveEvent(event)