custom port shape drawing #160

Author: jchanvfx

NodeGraphQt/base/node.py

@@ -529,6 +529,12 @@ def __init__(self):
         self._view.text_item.editingFinished.connect(self.set_name)
 
     def draw(self, force=True):
+        """
+        Redraws the node in the scene.
+
+        Args:
+            force (bool): force redraw if not visible.
+        """
         if force:
             if not self.model.visible:
                 self._has_draw = False
@@ -544,14 +550,16 @@ def hide(self):
         Hide node.
         """
         super(BaseNode, self).hide()
-        [pipe.hide() for port in self._inputs + self._outputs for pipe in port.view.connected_pipes]
+        [pipe.hide() for port in self._inputs + self._outputs
+         for pipe in port.view.connected_pipes]
 
     def show(self):
         """
         Show node.
         """
         super(BaseNode, self).show()
-        [pipe.show() for port in self._inputs + self._outputs for pipe in port.view.connected_pipes]
+        [pipe.show() for port in self._inputs + self._outputs
+         for pipe in port.view.connected_pipes]
         self.draw(False)
 
     def update_model(self):
@@ -755,7 +763,7 @@ def add_checkbox(self, name, label='', text='', state=False, tab=None):
         self.view.add_widget(widget)
 
     def add_input(self, name='input', multi_input=False, display_name=True,
-                  color=None, data_type='None'):
+                  color=None, data_type='None', painter_func=None):
         """
         Add input :class:`Port` to node.
 
@@ -764,19 +772,26 @@ def add_input(self, name='input', multi_input=False, display_name=True,
             multi_input (bool): allow port to have more than one connection.
             display_name (bool): display the port name on the node.
             color (tuple): initial port color (r, g, b) ``0-255``.
-            data_type(str): port data type name.
+            data_type (str): port data type name.
+            painter_func (function): custom function to override the drawing
+                of the port shape see example:
 
         Returns:
             NodeGraphQt.Port: the created port object.
         """
         if name in self.inputs().keys():
             raise PortRegistrationError(
                 'port name "{}" already registered.'.format(name))
-        view = self.view.add_input(name, multi_input, display_name)
+
+        port_args = [name, multi_input, display_name]
+        if painter_func and not callable(painter_func):
+            port_args.append(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]
+
         port = Port(self, view)
         port.model.type_ = IN_PORT
         port.model.name = name
@@ -788,7 +803,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, data_type='None'):
+                   color=None, data_type='None', painter_func=None):
         """
         Add output :class:`Port` to node.
 
@@ -798,14 +813,21 @@ def add_output(self, name='output', multi_output=True, display_name=True,
             display_name (bool): display the port name on the node.
             color (tuple): initial port color (r, g, b) ``0-255``.
             data_type(str): port data type name.
+            painter_func (function): custom function to override the drawing
+                of the port shape see example:
 
         Returns:
             NodeGraphQt.Port: the created port object.
         """
         if name in self.outputs().keys():
             raise PortRegistrationError(
                 'port name "{}" already registered.'.format(name))
-        view = self.view.add_output(name, multi_output, display_name)
+
+        port_args = [name, multi_output, display_name]
+        if painter_func and callable(painter_func):
+            port_args.append(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]
@@ -914,8 +936,21 @@ def set_ports(self, port_data):
         Set node input and output ports.
 
         Args:
-            port_data(dict): {'input_ports':[{'name':...,'multi_connection':...,'display_name':...,'data_type':...}, ...],
-            "                 'output_ports':[{'name':...,'multi_connection':...,'display_name':...,'data_type':...}, ...]}
+            port_data(dict):
+                {
+                'input_ports':
+                    [{'name':...,
+                      'multi_connection':...,
+                      'display_name':...,
+                      'data_type':...
+                    }, ...],
+                'output_ports':
+                    [{'name':...,
+                      'multi_connection':...,
+                      'display_name':...,
+                      'data_type':...
+                    }, ...]
+                }
         """
         for port in self._inputs:
             self._view.delete_input(port.view)
@@ -927,11 +962,15 @@ def set_ports(self, port_data):
         self._outputs = []
         self._model.outputs = {}
         self._model.inputs = {}
-        [self.add_input(name=port['name'], multi_input=port['multi_connection'],
-                        display_name=port['display_name'], data_type=port['data_type'])
+        [self.add_input(name=port['name'],
+                        multi_input=port['multi_connection'],
+                        display_name=port['display_name'],
+                        data_type=port['data_type'])
          for port in port_data['input_ports']]
-        [self.add_output(name=port['name'], multi_output=port['multi_connection'],
-                         display_name=port['display_name'], data_type=port['data_type'])
+        [self.add_output(name=port['name'],
+                         multi_output=port['multi_connection'],
+                         display_name=port['display_name'],
+                         data_type=port['data_type'])
          for port in port_data['output_ports']]
         self.draw()
 
@@ -1216,6 +1255,7 @@ def create_from_nodes(self, nodes):
     def add_child(self, node):
         """
         Add a node to the sub graph.
+
         Args:
             node(NodeGraphQt.BaseNode).
         """
@@ -1224,6 +1264,7 @@ def add_child(self, node):
     def remove_child(self, node):
         """
         Remove a node from the sub graph.
+
         Args:
             node(NodeGraphQt.BaseNode).
         """

NodeGraphQt/base/port.py

@@ -1,9 +1,9 @@
 #!/usr/bin/python
 from .commands import (PortConnectedCmd,
-                             PortDisconnectedCmd,
-                             PortVisibleCmd,
-                             NodeInputConnectedCmd,
-                             NodeInputDisconnectedCmd)
+                       PortDisconnectedCmd,
+                       PortVisibleCmd,
+                       NodeInputConnectedCmd,
+                       NodeInputDisconnectedCmd)
 from .model import PortModel
 from ..constants import IN_PORT, OUT_PORT
 
@@ -221,7 +221,6 @@ def data_type(self):
     def data_type(self, data_type):
         self.__model.data_type = data_type
 
-
     @property
     def border_color(self):
         return self.__view.border_color