1# -*- coding: utf-8 -*-
2"""
3Simple example of subclassing GraphItem.
4"""
5
6import initExample ## Add path to library (just for examples; you do not need this)
7
8import pyqtgraph as pg
9from pyqtgraph.Qt import QtCore, QtGui
10import numpy as np
11
12# Enable antialiasing for prettier plots
13pg.setConfigOptions(antialias=True)
14
15w = pg.GraphicsLayoutWidget(show=True)
16w.setWindowTitle('pyqtgraph example: CustomGraphItem')
17v = w.addViewBox()
18v.setAspectLocked()
19
20class Graph(pg.GraphItem):
21    def __init__(self):
22        self.dragPoint = None
23        self.dragOffset = None
24        self.textItems = []
25        pg.GraphItem.__init__(self)
26        self.scatter.sigClicked.connect(self.clicked)
27
28    def setData(self, **kwds):
29        self.text = kwds.pop('text', [])
30        self.data = kwds
31        if 'pos' in self.data:
32            npts = self.data['pos'].shape[0]
33            self.data['data'] = np.empty(npts, dtype=[('index', int)])
34            self.data['data']['index'] = np.arange(npts)
35        self.setTexts(self.text)
36        self.updateGraph()
37
38    def setTexts(self, text):
39        for i in self.textItems:
40            i.scene().removeItem(i)
41        self.textItems = []
42        for t in text:
43            item = pg.TextItem(t)
44            self.textItems.append(item)
45            item.setParentItem(self)
46
47    def updateGraph(self):
48        pg.GraphItem.setData(self, **self.data)
49        for i,item in enumerate(self.textItems):
50            item.setPos(*self.data['pos'][i])
51
52
53    def mouseDragEvent(self, ev):
54        if ev.button() != QtCore.Qt.MouseButton.LeftButton:
55            ev.ignore()
56            return
57
58        if ev.isStart():
59            # We are already one step into the drag.
60            # Find the point(s) at the mouse cursor when the button was first
61            # pressed:
62            pos = ev.buttonDownPos()
63            pts = self.scatter.pointsAt(pos)
64            if len(pts) == 0:
65                ev.ignore()
66                return
67            self.dragPoint = pts[0]
68            ind = pts[0].data()[0]
69            self.dragOffset = self.data['pos'][ind] - pos
70        elif ev.isFinish():
71            self.dragPoint = None
72            return
73        else:
74            if self.dragPoint is None:
75                ev.ignore()
76                return
77
78        ind = self.dragPoint.data()[0]
79        self.data['pos'][ind] = ev.pos() + self.dragOffset
80        self.updateGraph()
81        ev.accept()
82
83    def clicked(self, pts):
84        print("clicked: %s" % pts)
85
86
87g = Graph()
88v.addItem(g)
89
90## Define positions of nodes
91pos = np.array([
92    [0,0],
93    [10,0],
94    [0,10],
95    [10,10],
96    [5,5],
97    [15,5]
98    ], dtype=float)
99
100## Define the set of connections in the graph
101adj = np.array([
102    [0,1],
103    [1,3],
104    [3,2],
105    [2,0],
106    [1,5],
107    [3,5],
108    ])
109
110## Define the symbol to use for each node (this is optional)
111symbols = ['o','o','o','o','t','+']
112
113## Define the line style for each connection (this is optional)
114lines = np.array([
115    (255,0,0,255,1),
116    (255,0,255,255,2),
117    (255,0,255,255,3),
118    (255,255,0,255,2),
119    (255,0,0,255,1),
120    (255,255,255,255,4),
121    ], dtype=[('red',np.ubyte),('green',np.ubyte),('blue',np.ubyte),('alpha',np.ubyte),('width',float)])
122
123## Define text to show next to each symbol
124texts = ["Point %d" % i for i in range(6)]
125
126## Update the graph
127g.setData(pos=pos, adj=adj, pen=lines, size=1, symbol=symbols, pxMode=False, text=texts)
128
129
130
131
132if __name__ == '__main__':
133    pg.exec()
134