From 237b8488371a7b68f224927c8b269ec6bbb2b41e Mon Sep 17 00:00:00 2001 From: Luke Campagnola Date: Fri, 15 Sep 2017 08:59:15 -0700 Subject: [PATCH] Allow binary operator nodes to select output type --- pyqtgraph/flowchart/library/Operators.py | 29 ++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/pyqtgraph/flowchart/library/Operators.py b/pyqtgraph/flowchart/library/Operators.py index 579d2cd2..596e8854 100644 --- a/pyqtgraph/flowchart/library/Operators.py +++ b/pyqtgraph/flowchart/library/Operators.py @@ -1,5 +1,7 @@ # -*- coding: utf-8 -*- from ..Node import Node +from .common import CtrlNode + class UniOpNode(Node): """Generic node for performing any operation like Out = In.fn()""" @@ -13,11 +15,22 @@ class UniOpNode(Node): def process(self, **args): return {'Out': getattr(args['In'], self.fn)()} -class BinOpNode(Node): +class BinOpNode(CtrlNode): """Generic node for performing any operation like A.fn(B)""" + + _dtypes = [ + 'float64', 'float32', 'float16', + 'int64', 'int32', 'int16', 'int8', + 'uint64', 'uint32', 'uint16', 'uint8' + ] + + uiTemplate = [ + ('outputType', 'combo', {'values': ['no change', 'input A', 'input B'] + _dtypes , 'index': 0}) + ] + def __init__(self, name, fn): self.fn = fn - Node.__init__(self, name, terminals={ + CtrlNode.__init__(self, name, terminals={ 'A': {'io': 'in'}, 'B': {'io': 'in'}, 'Out': {'io': 'out', 'bypass': 'A'} @@ -36,6 +49,18 @@ class BinOpNode(Node): out = fn(args['B']) if out is NotImplemented: raise Exception("Operation %s not implemented between %s and %s" % (fn, str(type(args['A'])), str(type(args['B'])))) + + # Coerce dtype if requested + typ = self.stateGroup.state()['outputType'] + if typ == 'no change': + pass + elif typ == 'input A': + out = out.astype(args['A'].dtype) + elif typ == 'input B': + out = out.astype(args['B'].dtype) + else: + out = out.astype(typ) + #print " ", fn, out return {'Out': out}