pyqtgraph/flowchart/Node.py

541 lines
18 KiB
Python
Raw Normal View History

# -*- coding: utf-8 -*-
from pyqtgraph.Qt import QtCore, QtGui
from pyqtgraph.graphicsItems.GraphicsObject import GraphicsObject
import pyqtgraph.functions as fn
from .Terminal import *
from pyqtgraph.pgcollections import OrderedDict
from pyqtgraph.debug import *
import numpy as np
from .eq import *
def strDict(d):
return dict([(str(k), v) for k, v in d.items()])
class Node(QtCore.QObject):
sigOutputChanged = QtCore.Signal(object) # self
sigClosed = QtCore.Signal(object)
sigRenamed = QtCore.Signal(object, object)
sigTerminalRenamed = QtCore.Signal(object, object) # term, oldName
sigTerminalAdded = QtCore.Signal(object, object) # self, term
sigTerminalRemoved = QtCore.Signal(object, object) # self, term
def __init__(self, name, terminals=None, allowAddInput=False, allowAddOutput=False, allowRemove=True):
QtCore.QObject.__init__(self)
self._name = name
self._bypass = False
self.bypassButton = None ## this will be set by the flowchart ctrl widget..
self._graphicsItem = None
self.terminals = OrderedDict()
self._inputs = OrderedDict()
self._outputs = OrderedDict()
self._allowAddInput = allowAddInput ## flags to allow the user to add/remove terminals
self._allowAddOutput = allowAddOutput
self._allowRemove = allowRemove
self.exception = None
if terminals is None:
return
for name, opts in terminals.items():
self.addTerminal(name, **opts)
def nextTerminalName(self, name):
"""Return an unused terminal name"""
name2 = name
i = 1
while name2 in self.terminals:
name2 = "%s.%d" % (name, i)
i += 1
return name2
def addInput(self, name="Input", **args):
#print "Node.addInput called."
return self.addTerminal(name, io='in', **args)
def addOutput(self, name="Output", **args):
return self.addTerminal(name, io='out', **args)
def removeTerminal(self, term):
## term may be a terminal or its name
if isinstance(term, Terminal):
name = term.name()
else:
name = term
term = self.terminals[name]
#print "remove", name
#term.disconnectAll()
term.close()
del self.terminals[name]
if name in self._inputs:
del self._inputs[name]
if name in self._outputs:
del self._outputs[name]
self.graphicsItem().updateTerminals()
self.sigTerminalRemoved.emit(self, term)
def terminalRenamed(self, term, oldName):
"""Called after a terminal has been renamed"""
newName = term.name()
for d in [self.terminals, self._inputs, self._outputs]:
if oldName not in d:
continue
d[newName] = d[oldName]
del d[oldName]
self.graphicsItem().updateTerminals()
self.sigTerminalRenamed.emit(term, oldName)
def addTerminal(self, name, **opts):
name = self.nextTerminalName(name)
term = Terminal(self, name, **opts)
self.terminals[name] = term
if term.isInput():
self._inputs[name] = term
elif term.isOutput():
self._outputs[name] = term
self.graphicsItem().updateTerminals()
self.sigTerminalAdded.emit(self, term)
return term
def inputs(self):
return self._inputs
def outputs(self):
return self._outputs
def process(self, **kargs):
"""Process data through this node. Each named argument supplies data to the corresponding terminal."""
return {}
def graphicsItem(self):
"""Return a (the?) graphicsitem for this node"""
#print "Node.graphicsItem called."
if self._graphicsItem is None:
#print "Creating NodeGraphicsItem..."
self._graphicsItem = NodeGraphicsItem(self)
#print "Node.graphicsItem is returning ", self._graphicsItem
return self._graphicsItem
def __getattr__(self, attr):
"""Return the terminal with the given name"""
if attr not in self.terminals:
raise AttributeError(attr)
else:
return self.terminals[attr]
def __getitem__(self, item):
return getattr(self, item)
def name(self):
return self._name
def rename(self, name):
oldName = self._name
self._name = name
#self.emit(QtCore.SIGNAL('renamed'), self, oldName)
self.sigRenamed.emit(self, oldName)
def dependentNodes(self):
"""Return the list of nodes which provide direct input to this node"""
nodes = set()
for t in self.inputs().values():
nodes |= set([i.node() for i in t.inputTerminals()])
return nodes
#return set([t.inputTerminals().node() for t in self.listInputs().itervalues()])
def __repr__(self):
return "<Node %s @%x>" % (self.name(), id(self))
def ctrlWidget(self):
return None
def bypass(self, byp):
self._bypass = byp
if self.bypassButton is not None:
self.bypassButton.setChecked(byp)
self.update()
def isBypassed(self):
return self._bypass
def setInput(self, **args):
"""Set the values on input terminals. For most nodes, this will happen automatically through Terminal.inputChanged.
This is normally only used for nodes with no connected inputs."""
changed = False
for k, v in args.items():
term = self._inputs[k]
oldVal = term.value()
if not eq(oldVal, v):
changed = True
term.setValue(v, process=False)
if changed and '_updatesHandled_' not in args:
self.update()
def inputValues(self):
vals = {}
for n, t in self.inputs().items():
vals[n] = t.value()
return vals
def outputValues(self):
vals = {}
for n, t in self.outputs().items():
vals[n] = t.value()
return vals
def connected(self, localTerm, remoteTerm):
"""Called whenever one of this node's terminals is connected elsewhere."""
pass
def disconnected(self, localTerm, remoteTerm):
"""Called whenever one of this node's terminals is connected elsewhere."""
pass
def update(self, signal=True):
"""Collect all input values, attempt to process new output values, and propagate downstream."""
vals = self.inputValues()
#print " inputs:", vals
try:
if self.isBypassed():
out = self.processBypassed(vals)
else:
out = self.process(**strDict(vals))
#print " output:", out
if out is not None:
if signal:
self.setOutput(**out)
else:
self.setOutputNoSignal(**out)
for n,t in self.inputs().items():
t.setValueAcceptable(True)
self.clearException()
except:
#printExc( "Exception while processing %s:" % self.name())
for n,t in self.outputs().items():
t.setValue(None)
self.setException(sys.exc_info())
if signal:
#self.emit(QtCore.SIGNAL('outputChanged'), self) ## triggers flowchart to propagate new data
self.sigOutputChanged.emit(self) ## triggers flowchart to propagate new data
def processBypassed(self, args):
result = {}
for term in list(self.outputs().values()):
byp = term.bypassValue()
if byp is None:
result[term.name()] = None
else:
result[term.name()] = args.get(byp, None)
return result
def setOutput(self, **vals):
self.setOutputNoSignal(**vals)
#self.emit(QtCore.SIGNAL('outputChanged'), self) ## triggers flowchart to propagate new data
self.sigOutputChanged.emit(self) ## triggers flowchart to propagate new data
def setOutputNoSignal(self, **vals):
for k, v in vals.items():
term = self.outputs()[k]
term.setValue(v)
#targets = term.connections()
#for t in targets: ## propagate downstream
#if t is term:
#continue
#t.inputChanged(term)
term.setValueAcceptable(True)
def setException(self, exc):
self.exception = exc
self.recolor()
def clearException(self):
self.setException(None)
def recolor(self):
if self.exception is None:
self.graphicsItem().setPen(QtGui.QPen(QtGui.QColor(0, 0, 0)))
else:
self.graphicsItem().setPen(QtGui.QPen(QtGui.QColor(150, 0, 0), 3))
def saveState(self):
pos = self.graphicsItem().pos()
state = {'pos': (pos.x(), pos.y()), 'bypass': self.isBypassed()}
termsEditable = self._allowAddInput | self._allowAddOutput
for term in self._inputs.values() + self._outputs.values():
termsEditable |= term._renamable | term._removable | term._multiable
if termsEditable:
state['terminals'] = self.saveTerminals()
return state
def restoreState(self, state):
pos = state.get('pos', (0,0))
self.graphicsItem().setPos(*pos)
self.bypass(state.get('bypass', False))
if 'terminals' in state:
self.restoreTerminals(state['terminals'])
def saveTerminals(self):
terms = OrderedDict()
for n, t in self.terminals.items():
terms[n] = (t.saveState())
return terms
def restoreTerminals(self, state):
for name in list(self.terminals.keys()):
if name not in state:
self.removeTerminal(name)
for name, opts in state.items():
if name in self.terminals:
term = self[name]
term.setOpts(**opts)
continue
try:
opts = strDict(opts)
self.addTerminal(name, **opts)
except:
printExc("Error restoring terminal %s (%s):" % (str(name), str(opts)))
def clearTerminals(self):
for t in self.terminals.values():
t.close()
self.terminals = OrderedDict()
self._inputs = OrderedDict()
self._outputs = OrderedDict()
def close(self):
"""Cleans up after the node--removes terminals, graphicsItem, widget"""
self.disconnectAll()
self.clearTerminals()
item = self.graphicsItem()
if item.scene() is not None:
item.scene().removeItem(item)
self._graphicsItem = None
w = self.ctrlWidget()
if w is not None:
w.setParent(None)
#self.emit(QtCore.SIGNAL('closed'), self)
self.sigClosed.emit(self)
def disconnectAll(self):
for t in self.terminals.values():
t.disconnectAll()
#class NodeGraphicsItem(QtGui.QGraphicsItem):
class NodeGraphicsItem(GraphicsObject):
def __init__(self, node):
#QtGui.QGraphicsItem.__init__(self)
GraphicsObject.__init__(self)
#QObjectWorkaround.__init__(self)
#self.shadow = QtGui.QGraphicsDropShadowEffect()
#self.shadow.setOffset(5,5)
#self.shadow.setBlurRadius(10)
#self.setGraphicsEffect(self.shadow)
self.pen = fn.mkPen(0,0,0)
self.selectPen = fn.mkPen(200,200,200,width=2)
self.brush = fn.mkBrush(200, 200, 200, 150)
self.hoverBrush = fn.mkBrush(200, 200, 200, 200)
self.selectBrush = fn.mkBrush(200, 200, 255, 200)
self.hovered = False
self.node = node
flags = self.ItemIsMovable | self.ItemIsSelectable | self.ItemIsFocusable |self.ItemSendsGeometryChanges
#flags = self.ItemIsFocusable |self.ItemSendsGeometryChanges
self.setFlags(flags)
self.bounds = QtCore.QRectF(0, 0, 100, 100)
self.nameItem = QtGui.QGraphicsTextItem(self.node.name(), self)
self.nameItem.setDefaultTextColor(QtGui.QColor(50, 50, 50))
self.nameItem.moveBy(self.bounds.width()/2. - self.nameItem.boundingRect().width()/2., 0)
self.nameItem.setTextInteractionFlags(QtCore.Qt.TextEditorInteraction)
self.updateTerminals()
#self.setZValue(10)
self.nameItem.focusOutEvent = self.labelFocusOut
self.nameItem.keyPressEvent = self.labelKeyPress
self.menu = None
self.buildMenu()
#self.node.sigTerminalRenamed.connect(self.updateActionMenu)
#def setZValue(self, z):
#for t, item in self.terminals.itervalues():
#item.setZValue(z+1)
#GraphicsObject.setZValue(self, z)
def labelFocusOut(self, ev):
QtGui.QGraphicsTextItem.focusOutEvent(self.nameItem, ev)
self.labelChanged()
def labelKeyPress(self, ev):
if ev.key() == QtCore.Qt.Key_Enter or ev.key() == QtCore.Qt.Key_Return:
self.labelChanged()
else:
QtGui.QGraphicsTextItem.keyPressEvent(self.nameItem, ev)
def labelChanged(self):
newName = str(self.nameItem.toPlainText())
if newName != self.node.name():
self.node.rename(newName)
### re-center the label
bounds = self.boundingRect()
self.nameItem.setPos(bounds.width()/2. - self.nameItem.boundingRect().width()/2., 0)
def setPen(self, pen):
self.pen = pen
self.update()
def setBrush(self, brush):
self.brush = brush
self.update()
def updateTerminals(self):
bounds = self.bounds
self.terminals = {}
inp = self.node.inputs()
dy = bounds.height() / (len(inp)+1)
y = dy
for i, t in inp.items():
item = t.graphicsItem()
item.setParentItem(self)
#item.setZValue(self.zValue()+1)
br = self.bounds
item.setAnchor(0, y)
self.terminals[i] = (t, item)
y += dy
out = self.node.outputs()
dy = bounds.height() / (len(out)+1)
y = dy
for i, t in out.items():
item = t.graphicsItem()
item.setParentItem(self)
item.setZValue(self.zValue())
br = self.bounds
item.setAnchor(bounds.width(), y)
self.terminals[i] = (t, item)
y += dy
#self.buildMenu()
def boundingRect(self):
return self.bounds.adjusted(-5, -5, 5, 5)
def paint(self, p, *args):
p.setPen(self.pen)
if self.isSelected():
p.setPen(self.selectPen)
p.setBrush(self.selectBrush)
else:
p.setPen(self.pen)
if self.hovered:
p.setBrush(self.hoverBrush)
else:
p.setBrush(self.brush)
p.drawRect(self.bounds)
def mousePressEvent(self, ev):
ev.ignore()
def mouseClickEvent(self, ev):
#print "Node.mouseClickEvent called."
if int(ev.button()) == int(QtCore.Qt.LeftButton):
ev.accept()
#print " ev.button: left"
sel = self.isSelected()
#ret = QtGui.QGraphicsItem.mousePressEvent(self, ev)
self.setSelected(True)
if not sel and self.isSelected():
#self.setBrush(QtGui.QBrush(QtGui.QColor(200, 200, 255)))
#self.emit(QtCore.SIGNAL('selected'))
#self.scene().selectionChanged.emit() ## for some reason this doesn't seem to be happening automatically
self.update()
#return ret
elif int(ev.button()) == int(QtCore.Qt.RightButton):
#print " ev.button: right"
ev.accept()
#pos = ev.screenPos()
self.raiseContextMenu(ev)
#self.menu.popup(QtCore.QPoint(pos.x(), pos.y()))
def mouseDragEvent(self, ev):
#print "Node.mouseDrag"
if ev.button() == QtCore.Qt.LeftButton:
ev.accept()
self.setPos(self.pos()+self.mapToParent(ev.pos())-self.mapToParent(ev.lastPos()))
def hoverEvent(self, ev):
if not ev.isExit() and ev.acceptClicks(QtCore.Qt.LeftButton):
ev.acceptDrags(QtCore.Qt.LeftButton)
self.hovered = True
else:
self.hovered = False
self.update()
def keyPressEvent(self, ev):
if ev.key() == QtCore.Qt.Key_Delete or ev.key() == QtCore.Qt.Key_Backspace:
ev.accept()
if not self.node._allowRemove:
return
self.node.close()
else:
ev.ignore()
def itemChange(self, change, val):
if change == self.ItemPositionHasChanged:
for k, t in self.terminals.items():
t[1].nodeMoved()
return GraphicsObject.itemChange(self, change, val)
def getMenu(self):
return self.menu
def getContextMenus(self, event):
return [self.menu]
def raiseContextMenu(self, ev):
menu = self.scene().addParentContextMenus(self, self.getMenu(), ev)
pos = ev.screenPos()
menu.popup(QtCore.QPoint(pos.x(), pos.y()))
def buildMenu(self):
self.menu = QtGui.QMenu()
self.menu.setTitle("Node")
a = self.menu.addAction("Add input", self.addInputFromMenu)
if not self.node._allowAddInput:
a.setEnabled(False)
a = self.menu.addAction("Add output", self.addOutputFromMenu)
if not self.node._allowAddOutput:
a.setEnabled(False)
a = self.menu.addAction("Remove node", self.node.close)
if not self.node._allowRemove:
a.setEnabled(False)
def addInputFromMenu(self): ## called when add input is clicked in context menu
self.node.addInput(renamable=True, removable=True, multiable=True)
def addOutputFromMenu(self): ## called when add output is clicked in context menu
self.node.addOutput(renamable=True, removable=True, multiable=False)