1#!/usr/bin/env python 2 3 4############################################################################# 5## 6## Copyright (C) 2013 Riverbank Computing Limited. 7## Copyright (C) 2010 Nokia Corporation and/or its subsidiary(-ies). 8## All rights reserved. 9## 10## This file is part of the examples of PyQt. 11## 12## $QT_BEGIN_LICENSE:BSD$ 13## You may use this file under the terms of the BSD license as follows: 14## 15## "Redistribution and use in source and binary forms, with or without 16## modification, are permitted provided that the following conditions are 17## met: 18## * Redistributions of source code must retain the above copyright 19## notice, this list of conditions and the following disclaimer. 20## * Redistributions in binary form must reproduce the above copyright 21## notice, this list of conditions and the following disclaimer in 22## the documentation and/or other materials provided with the 23## distribution. 24## * Neither the name of Nokia Corporation and its Subsidiary(-ies) nor 25## the names of its contributors may be used to endorse or promote 26## products derived from this software without specific prior written 27## permission. 28## 29## THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 30## "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 31## LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 32## A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 33## OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 34## SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 35## LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 36## DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 37## THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 38## (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 39## OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE." 40## $QT_END_LICENSE$ 41## 42############################################################################# 43 44 45import math 46 47from PyQt5.QtCore import (qAbs, QLineF, QPointF, qrand, QRectF, QSizeF, qsrand, 48 Qt, QTime) 49from PyQt5.QtGui import (QBrush, QColor, QLinearGradient, QPainter, 50 QPainterPath, QPen, QPolygonF, QRadialGradient) 51from PyQt5.QtWidgets import (QApplication, QGraphicsItem, QGraphicsScene, 52 QGraphicsView, QStyle) 53 54 55class Edge(QGraphicsItem): 56 Pi = math.pi 57 TwoPi = 2.0 * Pi 58 59 Type = QGraphicsItem.UserType + 2 60 61 def __init__(self, sourceNode, destNode): 62 super(Edge, self).__init__() 63 64 self.arrowSize = 10.0 65 self.sourcePoint = QPointF() 66 self.destPoint = QPointF() 67 68 self.setAcceptedMouseButtons(Qt.NoButton) 69 self.source = sourceNode 70 self.dest = destNode 71 self.source.addEdge(self) 72 self.dest.addEdge(self) 73 self.adjust() 74 75 def type(self): 76 return Edge.Type 77 78 def sourceNode(self): 79 return self.source 80 81 def setSourceNode(self, node): 82 self.source = node 83 self.adjust() 84 85 def destNode(self): 86 return self.dest 87 88 def setDestNode(self, node): 89 self.dest = node 90 self.adjust() 91 92 def adjust(self): 93 if not self.source or not self.dest: 94 return 95 96 line = QLineF(self.mapFromItem(self.source, 0, 0), 97 self.mapFromItem(self.dest, 0, 0)) 98 length = line.length() 99 100 self.prepareGeometryChange() 101 102 if length > 20.0: 103 edgeOffset = QPointF((line.dx() * 10) / length, 104 (line.dy() * 10) / length) 105 106 self.sourcePoint = line.p1() + edgeOffset 107 self.destPoint = line.p2() - edgeOffset 108 else: 109 self.sourcePoint = line.p1() 110 self.destPoint = line.p1() 111 112 def boundingRect(self): 113 if not self.source or not self.dest: 114 return QRectF() 115 116 penWidth = 1.0 117 extra = (penWidth + self.arrowSize) / 2.0 118 119 return QRectF(self.sourcePoint, 120 QSizeF(self.destPoint.x() - self.sourcePoint.x(), 121 self.destPoint.y() - self.sourcePoint.y())).normalized().adjusted(-extra, -extra, extra, extra) 122 123 def paint(self, painter, option, widget): 124 if not self.source or not self.dest: 125 return 126 127 # Draw the line itself. 128 line = QLineF(self.sourcePoint, self.destPoint) 129 130 if line.length() == 0.0: 131 return 132 133 painter.setPen(QPen(Qt.black, 1, Qt.SolidLine, Qt.RoundCap, 134 Qt.RoundJoin)) 135 painter.drawLine(line) 136 137 # Draw the arrows if there's enough room. 138 angle = math.acos(line.dx() / line.length()) 139 if line.dy() >= 0: 140 angle = Edge.TwoPi - angle 141 142 sourceArrowP1 = self.sourcePoint + QPointF(math.sin(angle + Edge.Pi / 3) * self.arrowSize, 143 math.cos(angle + Edge.Pi / 3) * self.arrowSize) 144 sourceArrowP2 = self.sourcePoint + QPointF(math.sin(angle + Edge.Pi - Edge.Pi / 3) * self.arrowSize, 145 math.cos(angle + Edge.Pi - Edge.Pi / 3) * self.arrowSize); 146 destArrowP1 = self.destPoint + QPointF(math.sin(angle - Edge.Pi / 3) * self.arrowSize, 147 math.cos(angle - Edge.Pi / 3) * self.arrowSize) 148 destArrowP2 = self.destPoint + QPointF(math.sin(angle - Edge.Pi + Edge.Pi / 3) * self.arrowSize, 149 math.cos(angle - Edge.Pi + Edge.Pi / 3) * self.arrowSize) 150 151 painter.setBrush(Qt.black) 152 painter.drawPolygon(QPolygonF([line.p1(), sourceArrowP1, sourceArrowP2])) 153 painter.drawPolygon(QPolygonF([line.p2(), destArrowP1, destArrowP2])) 154 155 156class Node(QGraphicsItem): 157 Type = QGraphicsItem.UserType + 1 158 159 def __init__(self, graphWidget): 160 super(Node, self).__init__() 161 162 self.graph = graphWidget 163 self.edgeList = [] 164 self.newPos = QPointF() 165 166 self.setFlag(QGraphicsItem.ItemIsMovable) 167 self.setFlag(QGraphicsItem.ItemSendsGeometryChanges) 168 self.setCacheMode(QGraphicsItem.DeviceCoordinateCache) 169 self.setZValue(1) 170 171 def type(self): 172 return Node.Type 173 174 def addEdge(self, edge): 175 self.edgeList.append(edge) 176 edge.adjust() 177 178 def edges(self): 179 return self.edgeList 180 181 def calculateForces(self): 182 if not self.scene() or self.scene().mouseGrabberItem() is self: 183 self.newPos = self.pos() 184 return 185 186 # Sum up all forces pushing this item away. 187 xvel = 0.0 188 yvel = 0.0 189 for item in self.scene().items(): 190 if not isinstance(item, Node): 191 continue 192 193 line = QLineF(self.mapFromItem(item, 0, 0), QPointF(0, 0)) 194 dx = line.dx() 195 dy = line.dy() 196 l = 2.0 * (dx * dx + dy * dy) 197 if l > 0: 198 xvel += (dx * 150.0) / l 199 yvel += (dy * 150.0) / l 200 201 # Now subtract all forces pulling items together. 202 weight = (len(self.edgeList) + 1) * 10.0 203 for edge in self.edgeList: 204 if edge.sourceNode() is self: 205 pos = self.mapFromItem(edge.destNode(), 0, 0) 206 else: 207 pos = self.mapFromItem(edge.sourceNode(), 0, 0) 208 xvel += pos.x() / weight 209 yvel += pos.y() / weight 210 211 if qAbs(xvel) < 0.1 and qAbs(yvel) < 0.1: 212 xvel = yvel = 0.0 213 214 sceneRect = self.scene().sceneRect() 215 self.newPos = self.pos() + QPointF(xvel, yvel) 216 self.newPos.setX(min(max(self.newPos.x(), sceneRect.left() + 10), sceneRect.right() - 10)) 217 self.newPos.setY(min(max(self.newPos.y(), sceneRect.top() + 10), sceneRect.bottom() - 10)) 218 219 def advance(self): 220 if self.newPos == self.pos(): 221 return False 222 223 self.setPos(self.newPos) 224 return True 225 226 def boundingRect(self): 227 adjust = 2.0 228 return QRectF(-10 - adjust, -10 - adjust, 23 + adjust, 23 + adjust) 229 230 def shape(self): 231 path = QPainterPath() 232 path.addEllipse(-10, -10, 20, 20) 233 return path 234 235 def paint(self, painter, option, widget): 236 painter.setPen(Qt.NoPen) 237 painter.setBrush(Qt.darkGray) 238 painter.drawEllipse(-7, -7, 20, 20) 239 240 gradient = QRadialGradient(-3, -3, 10) 241 if option.state & QStyle.State_Sunken: 242 gradient.setCenter(3, 3) 243 gradient.setFocalPoint(3, 3) 244 gradient.setColorAt(1, QColor(Qt.yellow).lighter(120)) 245 gradient.setColorAt(0, QColor(Qt.darkYellow).lighter(120)) 246 else: 247 gradient.setColorAt(0, Qt.yellow) 248 gradient.setColorAt(1, Qt.darkYellow) 249 250 painter.setBrush(QBrush(gradient)) 251 painter.setPen(QPen(Qt.black, 0)) 252 painter.drawEllipse(-10, -10, 20, 20) 253 254 def itemChange(self, change, value): 255 if change == QGraphicsItem.ItemPositionHasChanged: 256 for edge in self.edgeList: 257 edge.adjust() 258 self.graph.itemMoved() 259 260 return super(Node, self).itemChange(change, value) 261 262 def mousePressEvent(self, event): 263 self.update() 264 super(Node, self).mousePressEvent(event) 265 266 def mouseReleaseEvent(self, event): 267 self.update() 268 super(Node, self).mouseReleaseEvent(event) 269 270 271class GraphWidget(QGraphicsView): 272 def __init__(self): 273 super(GraphWidget, self).__init__() 274 275 self.timerId = 0 276 277 scene = QGraphicsScene(self) 278 scene.setItemIndexMethod(QGraphicsScene.NoIndex) 279 scene.setSceneRect(-200, -200, 400, 400) 280 self.setScene(scene) 281 self.setCacheMode(QGraphicsView.CacheBackground) 282 self.setViewportUpdateMode(QGraphicsView.BoundingRectViewportUpdate) 283 self.setRenderHint(QPainter.Antialiasing) 284 self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse) 285 self.setResizeAnchor(QGraphicsView.AnchorViewCenter) 286 287 node1 = Node(self) 288 node2 = Node(self) 289 node3 = Node(self) 290 node4 = Node(self) 291 self.centerNode = Node(self) 292 node6 = Node(self) 293 node7 = Node(self) 294 node8 = Node(self) 295 node9 = Node(self) 296 scene.addItem(node1) 297 scene.addItem(node2) 298 scene.addItem(node3) 299 scene.addItem(node4) 300 scene.addItem(self.centerNode) 301 scene.addItem(node6) 302 scene.addItem(node7) 303 scene.addItem(node8) 304 scene.addItem(node9) 305 scene.addItem(Edge(node1, node2)) 306 scene.addItem(Edge(node2, node3)) 307 scene.addItem(Edge(node2, self.centerNode)) 308 scene.addItem(Edge(node3, node6)) 309 scene.addItem(Edge(node4, node1)) 310 scene.addItem(Edge(node4, self.centerNode)) 311 scene.addItem(Edge(self.centerNode, node6)) 312 scene.addItem(Edge(self.centerNode, node8)) 313 scene.addItem(Edge(node6, node9)) 314 scene.addItem(Edge(node7, node4)) 315 scene.addItem(Edge(node8, node7)) 316 scene.addItem(Edge(node9, node8)) 317 318 node1.setPos(-50, -50) 319 node2.setPos(0, -50) 320 node3.setPos(50, -50) 321 node4.setPos(-50, 0) 322 self.centerNode.setPos(0, 0) 323 node6.setPos(50, 0) 324 node7.setPos(-50, 50) 325 node8.setPos(0, 50) 326 node9.setPos(50, 50) 327 328 self.scale(0.8, 0.8) 329 self.setMinimumSize(400, 400) 330 self.setWindowTitle("Elastic Nodes") 331 332 def itemMoved(self): 333 if not self.timerId: 334 self.timerId = self.startTimer(1000 / 25) 335 336 def keyPressEvent(self, event): 337 key = event.key() 338 339 if key == Qt.Key_Up: 340 self.centerNode.moveBy(0, -20) 341 elif key == Qt.Key_Down: 342 self.centerNode.moveBy(0, 20) 343 elif key == Qt.Key_Left: 344 self.centerNode.moveBy(-20, 0) 345 elif key == Qt.Key_Right: 346 self.centerNode.moveBy(20, 0) 347 elif key == Qt.Key_Plus: 348 self.scaleView(1.2) 349 elif key == Qt.Key_Minus: 350 self.scaleView(1 / 1.2) 351 elif key == Qt.Key_Space or key == Qt.Key_Enter: 352 for item in self.scene().items(): 353 if isinstance(item, Node): 354 item.setPos(-150 + qrand() % 300, -150 + qrand() % 300) 355 else: 356 super(GraphWidget, self).keyPressEvent(event) 357 358 def timerEvent(self, event): 359 nodes = [item for item in self.scene().items() if isinstance(item, Node)] 360 361 for node in nodes: 362 node.calculateForces() 363 364 itemsMoved = False 365 for node in nodes: 366 if node.advance(): 367 itemsMoved = True 368 369 if not itemsMoved: 370 self.killTimer(self.timerId) 371 self.timerId = 0 372 373 def wheelEvent(self, event): 374 self.scaleView(math.pow(2.0, -event.angleDelta().y() / 240.0)) 375 376 def drawBackground(self, painter, rect): 377 # Shadow. 378 sceneRect = self.sceneRect() 379 rightShadow = QRectF(sceneRect.right(), sceneRect.top() + 5, 5, 380 sceneRect.height()) 381 bottomShadow = QRectF(sceneRect.left() + 5, sceneRect.bottom(), 382 sceneRect.width(), 5) 383 if rightShadow.intersects(rect) or rightShadow.contains(rect): 384 painter.fillRect(rightShadow, Qt.darkGray) 385 if bottomShadow.intersects(rect) or bottomShadow.contains(rect): 386 painter.fillRect(bottomShadow, Qt.darkGray) 387 388 # Fill. 389 gradient = QLinearGradient(sceneRect.topLeft(), sceneRect.bottomRight()) 390 gradient.setColorAt(0, Qt.white) 391 gradient.setColorAt(1, Qt.lightGray) 392 painter.fillRect(rect.intersected(sceneRect), QBrush(gradient)) 393 painter.setBrush(Qt.NoBrush) 394 painter.drawRect(sceneRect) 395 396 # Text. 397 textRect = QRectF(sceneRect.left() + 4, sceneRect.top() + 4, 398 sceneRect.width() - 4, sceneRect.height() - 4) 399 message = "Click and drag the nodes around, and zoom with the " \ 400 "mouse wheel or the '+' and '-' keys" 401 402 font = painter.font() 403 font.setBold(True) 404 font.setPointSize(14) 405 painter.setFont(font) 406 painter.setPen(Qt.lightGray) 407 painter.drawText(textRect.translated(2, 2), message) 408 painter.setPen(Qt.black) 409 painter.drawText(textRect, message) 410 411 def scaleView(self, scaleFactor): 412 factor = self.transform().scale(scaleFactor, scaleFactor).mapRect(QRectF(0, 0, 1, 1)).width() 413 414 if factor < 0.07 or factor > 100: 415 return 416 417 self.scale(scaleFactor, scaleFactor) 418 419 420if __name__ == '__main__': 421 422 import sys 423 424 app = QApplication(sys.argv) 425 qsrand(QTime(0,0,0).secsTo(QTime.currentTime())) 426 427 widget = GraphWidget() 428 widget.show() 429 430 sys.exit(app.exec_()) 431