add Port types to Ports to adress

NodeGraphQt/base/model.py

@@ -10,14 +10,53 @@
 from NodeGraphQt.errors import NodePropertyError
 
 
+class PortDatatypeModel(object):
+    """
+    Data dump for a port datatype object.
+    """
+
+    def __init__(self, name="", color=None, painter_func=None):
+        self._name = name
+        self._color = color
+        self._painter_func = painter_func
+
+    @property
+    def name(self):
+        """Name of the datatype.
+
+        Returns:
+            str: name of the Datatype.
+        """
+        return self._name
+
+    @property
+    def color(self):
+        """Color of the datatype to override the default color.
+
+        Returns:
+            tuple: RGB color tuple.
+        """
+        return self._color
+
+    @property
+    def painter_func(self):
+        """Painter function for the datatype to override the default painter.
+
+        Returns:
+            callable: custom painter function.
+        """
+        return self._painter_func
+
+
 class PortModel(object):
     """
     Data dump for a port object.
     """
 
-    def __init__(self, node):
+    def __init__(self, node, data_type=None):
         self.node = node
         self.type_ = ''
+        self.data_type = data_type
         self.name = 'port'
         self.display_name = True
         self.multi_connection = False

NodeGraphQt/base/port.py

@@ -29,11 +29,12 @@ class Port(object):
     Args:
         node (NodeGraphQt.NodeObject): parent node.
         port (PortItem): graphic item used for drawing.
+        data_type (PortDatatypeModel): data type of the port.
     """
 
-    def __init__(self, node, port):
+    def __init__(self, node, port, data_type=None):
         self.__view = port
-        self.__model = PortModel(node)
+        self.__model = PortModel(node,data_type)
 
     def __repr__(self):
         port = str(self.__class__.__name__)
@@ -232,6 +233,8 @@ def connect_to(self, port=None, push_undo=True, emit_signal=True):
                 'Can\'t connect port because "{}" is locked.'.format(name))
 
         # validate accept connection.
+        if self.model.data_type != port.model.data_type:
+            return
         node_type = self.node().type_
         accepted_types = port.accepted_port_types().get(node_type)
         if accepted_types:

NodeGraphQt/nodes/base_node.py

@@ -375,7 +375,7 @@ def show_widget(self, name, push_undo=True):
             undo_cmd.redo()
 
     def add_input(self, name='input', multi_input=False, display_name=True,
-                  color=None, locked=False, painter_func=None):
+                  color=None, locked=False, painter_func=None, data_type=None):
         """
         Add input :class:`Port` to node.
 
@@ -390,6 +390,7 @@ def add_input(self, name='input', multi_input=False, display_name=True,
             locked (bool): locked state see :meth:`Port.set_locked`
             painter_func (function or None): custom function to override the drawing
                 of the port shape see example: :ref:`Creating Custom Shapes`
+            data_type (PortDatatypeModel): data type for the port.
 
         Returns:
             NodeGraphQt.Port: the created port object.
@@ -401,13 +402,17 @@ def add_input(self, name='input', multi_input=False, display_name=True,
         port_args = [name, multi_input, display_name, locked]
         if painter_func and callable(painter_func):
             port_args.append(painter_func)
+        elif data_type and data_type.painter_func:
+            port_args.append(data_type.painter_func)
         view = self.view.add_input(*port_args)
 
         if color:
             view.color = color
             view.border_color = [min([255, max([0, i + 80])]) for i in color]
+        elif data_type and data_type.color:
+            view.color = data_type.color
 
-        port = Port(self, view)
+        port = Port(self, view, data_type)
         port.model.type_ = PortTypeEnum.IN.value
         port.model.name = name
         port.model.display_name = display_name
@@ -418,7 +423,7 @@ def add_input(self, name='input', multi_input=False, display_name=True,
         return port
 
     def add_output(self, name='output', multi_output=True, display_name=True,
-                   color=None, locked=False, painter_func=None):
+                   color=None, locked=False, painter_func=None, data_type=None):
         """
         Add output :class:`Port` to node.
 
@@ -433,6 +438,7 @@ def add_output(self, name='output', multi_output=True, display_name=True,
             locked (bool): locked state see :meth:`Port.set_locked`
             painter_func (function or None): custom function to override the drawing
                 of the port shape see example: :ref:`Creating Custom Shapes`
+            data_type (PortDatatypeModel): data type for the port.
 
         Returns:
             NodeGraphQt.Port: the created port object.
@@ -444,12 +450,16 @@ def add_output(self, name='output', multi_output=True, display_name=True,
         port_args = [name, multi_output, display_name, locked]
         if painter_func and callable(painter_func):
             port_args.append(painter_func)
+        elif data_type and data_type.painter_func:
+            port_args.append(data_type.painter_func)
         view = self.view.add_output(*port_args)
 
         if color:
             view.color = color
             view.border_color = [min([255, max([0, i + 80])]) for i in color]
-        port = Port(self, view)
+        elif data_type and data_type.color:
+            view.color = data_type.color
+        port = Port(self, view, data_type)
         port.model.type_ = PortTypeEnum.OUT.value
         port.model.name = name
         port.model.display_name = display_name

examples/nodes/custom_ports_node.py

@@ -2,6 +2,7 @@
 from Qt import QtCore, QtGui
 
 from NodeGraphQt import BaseNode
+from NodeGraphQt.base.model import PortDatatypeModel
 
 
 def draw_triangle_port(painter, rect, info):
@@ -110,12 +111,14 @@ class CustomPortsNode(BaseNode):
 
     # set the initial default node name.
     NODE_NAME = 'node'
+    int_port = PortDatatypeModel("int", (200, 10, 0), draw_triangle_port)
 
     def __init__(self):
         super(CustomPortsNode, self).__init__()
 
         # create input and output port.
         self.add_input('in', color=(200, 10, 0))
+        self.add_input('triangel_in', data_type=self.int_port)
         self.add_output('default')
         self.add_output('square', painter_func=draw_square_port)
-        self.add_output('triangle', painter_func=draw_triangle_port)
+        self.add_output('triangle_out', data_type=self.int_port)