group node tweaks

jchanvfx · jchanvfx · commit 9a1e3bf926fb · 2022-03-24T01:13:07.000+13:00
diff --git a/NodeGraphQt/__init__.py b/NodeGraphQt/__init__.py
@@ -76,7 +76,7 @@ def __init__(self):
 from .pkg_info import __license__ as LICENSE
 
 # node graph
-from .base.graph import NodeGraph
+from .base.graph import NodeGraph, SubGraph
 from .base.menu import NodesMenu, NodeGraphMenu, NodeGraphCommand
 
 # nodes & ports
@@ -109,6 +109,7 @@ def __init__(self):
     'NodesMenu',
     'Port',
     'PropertiesBinWidget',
+    'SubGraph',
     'VERSION',
     'constants',
     'custom_widgets'
diff --git a/NodeGraphQt/base/graph.py b/NodeGraphQt/base/graph.py
@@ -1947,18 +1947,30 @@ def _build_port_nodes(self):
                 input_node.add_output(port.name())
                 input_nodes[port.name()] = input_node
                 self.add_node(input_node, selected=False, push_undo=False)
+                x, y = input_node.pos()
+                if NODE_LAYOUT_DIRECTION is NODE_LAYOUT_HORIZONTAL:
+                    x -= 100
+                elif NODE_LAYOUT_DIRECTION is NODE_LAYOUT_VERTICAL:
+                    y -= 100
+                input_node.set_property('pos', [x, y], push_undo=False)
 
         # build the parent output port nodes.
         output_nodes = {n.name(): n for n in
                         self.get_nodes_by_type(PortOutputNode.type_)}
         for port in self.node.output_ports():
             if port.name() not in output_nodes:
-                output_port = PortOutputNode(parent_port=port)
-                output_port.NODE_NAME = port.name()
-                output_port.model.set_property('name', port.name())
-                output_port.add_input(port.name())
-                output_nodes[port.name()] = output_port
-                self.add_node(output_port, selected=False, push_undo=False)
+                output_node = PortOutputNode(parent_port=port)
+                output_node.NODE_NAME = port.name()
+                output_node.model.set_property('name', port.name())
+                output_node.add_input(port.name())
+                output_nodes[port.name()] = output_node
+                self.add_node(output_node, selected=False, push_undo=False)
+                x, y = output_node.pos()
+                if NODE_LAYOUT_DIRECTION is NODE_LAYOUT_HORIZONTAL:
+                    x += 100
+                elif NODE_LAYOUT_DIRECTION is NODE_LAYOUT_VERTICAL:
+                    y += 100
+                output_node.set_property('pos', [x, y], push_undo=False)
 
         return input_nodes, output_nodes
 
diff --git a/NodeGraphQt/nodes/group_node.py b/NodeGraphQt/nodes/group_node.py
@@ -3,6 +3,7 @@
                                    NODE_LAYOUT_HORIZONTAL)
 
 from NodeGraphQt.nodes.base_node import BaseNode
+from NodeGraphQt.nodes.port_node import PortInputNode, PortOutputNode
 from NodeGraphQt.qgraphics.node_group import (GroupNodeItem,
                                               GroupNodeVerticalItem)
 
@@ -14,6 +15,9 @@ class GroupNode(BaseNode):
     inside of it.
 
     **Inherited from:** :class:`NodeGraphQt.BaseNode`
+
+    .. image:: ../_images/group_node.png
+        :width: 250px
     """
 
     NODE_NAME = 'Group'
@@ -35,11 +39,14 @@ def is_expanded(self):
         Returns:
             bool: true if the node is expanded.
         """
+        if not self.graph:
+            return False
         return bool(self.id in self.graph.sub_graphs)
 
     def get_sub_graph(self):
         """
-        Returns the initialized sub node graph controller to the group node.
+        Returns the sub graph controller to the group node if initialized
+        or returns None.
 
         Returns:
             NodeGraphQt.SubGraph or None: sub graph controller.
@@ -84,3 +91,71 @@ def collapse(self):
             :meth:`SubGraph.collapse_group_node`.
         """
         self.graph.collapse_group_node(self)
+
+    def add_input(self, name='input', multi_input=False, display_name=True,
+                  color=None, locked=False, painter_func=None):
+        port = super(GroupNode, self).add_input(
+            name=name,
+            multi_input=multi_input,
+            display_name=display_name,
+            color=color,
+            locked=locked,
+            painter_func=painter_func
+        )
+        if self.is_expanded:
+            input_node = PortInputNode(parent_port=port)
+            input_node.NODE_NAME = port.name()
+            input_node.model.set_property('name', port.name())
+            input_node.add_output(port.name())
+            sub_graph = self.get_sub_graph()
+            sub_graph.add_node(input_node, selected=False, push_undo=False)
+
+        return port
+
+    def add_output(self, name='output', multi_output=True, display_name=True,
+                   color=None, locked=False, painter_func=None):
+        port = super(GroupNode, self).add_output(
+            name=name,
+            multi_output=multi_output,
+            display_name=display_name,
+            color=color,
+            locked=locked,
+            painter_func=painter_func
+        )
+        if self.is_expanded:
+            output_port = PortOutputNode(parent_port=port)
+            output_port.NODE_NAME = port.name()
+            output_port.model.set_property('name', port.name())
+            output_port.add_input(port.name())
+            sub_graph = self.get_sub_graph()
+            sub_graph.add_node(output_port, selected=False, push_undo=False)
+
+        return port
+
+    def delete_input(self, port):
+        if type(port) in [int, str]:
+            port = self.get_output(port)
+            if port is None:
+                return
+
+        if self.is_expanded:
+            sub_graph = self.get_sub_graph()
+            port_node = sub_graph.get_node_by_port(port)
+            if port_node:
+                sub_graph.remove_node(port_node, push_undo=False)
+
+        super(GroupNode, self).delete_input(port)
+
+    def delete_output(self, port):
+        if type(port) in [int, str]:
+            port = self.get_output(port)
+            if port is None:
+                return
+
+        if self.is_expanded:
+            sub_graph = self.get_sub_graph()
+            port_node = sub_graph.get_node_by_port(port)
+            if port_node:
+                sub_graph.remove_node(port_node, push_undo=False)
+
+        super(GroupNode, self).delete_output(port)
diff --git a/NodeGraphQt/nodes/port_node.py b/NodeGraphQt/nodes/port_node.py
@@ -2,6 +2,7 @@
 from NodeGraphQt.constants import (NODE_LAYOUT_VERTICAL,
                                    NODE_LAYOUT_HORIZONTAL)
 
+from NodeGraphQt.errors import PortRegistrationError
 from NodeGraphQt.nodes.base_node import BaseNode
 from NodeGraphQt.qgraphics.node_port_in import (PortInputNodeItem,
                                                 PortInputNodeVerticalItem)
@@ -16,6 +17,10 @@ class PortInputNode(BaseNode):
     :class:`NodeGraphQt.SubGraph`.
 
     **Inherited from:** :class:`NodeGraphQt.BaseNode`
+
+    .. image:: ../_images/port_in_node.png
+        :width: 150px
+
     """
 
     NODE_NAME = 'InputPort'
@@ -40,12 +45,21 @@ def parent_port(self):
 
     def add_input(self, name='input', multi_input=False, display_name=True,
                   color=None, locked=False, painter_func=None):
-        return
+        """
+        This is not available for the `PortInputNode` class.
+        """
+        raise PortRegistrationError(
+            '"{}.add_input()" is not available for {}.'
+            .format(self.__class__.__name__, self)
+        )
 
     def add_output(self, name='output', multi_output=True, display_name=True,
                    color=None, locked=False, painter_func=None):
         if self._outputs:
-            return
+            raise PortRegistrationError(
+                '"{}.add_output()" only ONE output is allowed for this node.'
+                .format(self.__class__.__name__, self)
+            )
         super(PortInputNode, self).add_output(
             name=name,
             multi_output=multi_output,
@@ -63,6 +77,10 @@ class PortOutputNode(BaseNode):
     :class:`NodeGraphQt.SubGraph`.
 
     **Inherited from:** :class:`NodeGraphQt.BaseNode`
+
+    .. image:: ../_images/port_out_node.png
+        :width: 150px
+
     """
 
     NODE_NAME = 'OutputPort'
@@ -88,7 +106,10 @@ def parent_port(self):
     def add_input(self, name='input', multi_input=False, display_name=True,
                   color=None, locked=False, painter_func=None):
         if self._inputs:
-            return
+            raise PortRegistrationError(
+                '"{}.add_input()" only ONE input is allowed for this node.'
+                .format(self.__class__.__name__, self)
+            )
         super(PortOutputNode, self).add_input(
             name=name,
             multi_input=multi_input,
@@ -100,4 +121,10 @@ def add_input(self, name='input', multi_input=False, display_name=True,
 
     def add_output(self, name='output', multi_output=True, display_name=True,
                    color=None, locked=False, painter_func=None):
-        return
+        """
+        This is not available for the `PortOutputNode` class.
+        """
+        raise PortRegistrationError(
+            '"{}.add_output()" is not available for {}.'
+            .format(self.__class__.__name__, self)
+        )
diff --git a/NodeGraphQt/qgraphics/node_port_in.py b/NodeGraphQt/qgraphics/node_port_in.py
@@ -56,6 +56,9 @@ def paint(self, painter, option, widget):
             text_rect.height()
         )
 
+        painter.setBrush(QtGui.QColor(255, 255, 255, 20))
+        painter.drawRoundedRect(rect, 20, 20)
+
         painter.setBrush(QtGui.QColor(0, 0, 0, 100))
         painter.drawRoundedRect(text_rect, 3, 3)
 
@@ -81,8 +84,10 @@ def paint(self, painter, option, widget):
         painter.setPen(pen)
         painter.drawPolygon(poly)
 
-        edge_rect = QtCore.QRectF(
-            rect.width() - (size * 1.8), rect.center().y() - 15, 4, 30)
+        edge_size = 30
+        edge_rect = QtCore.QRectF(rect.width() - (size * 1.7),
+                                  rect.center().y() - (edge_size / 2),
+                                  4, edge_size)
         painter.drawRect(edge_rect)
 
         painter.restore()
diff --git a/NodeGraphQt/qgraphics/node_port_out.py b/NodeGraphQt/qgraphics/node_port_out.py
@@ -56,6 +56,9 @@ def paint(self, painter, option, widget):
             text_rect.height()
         )
 
+        painter.setBrush(QtGui.QColor(255, 255, 255, 20))
+        painter.drawRoundedRect(rect, 20, 20)
+
         painter.setBrush(QtGui.QColor(0, 0, 0, 100))
         painter.drawRoundedRect(text_rect, 3, 3)
 
@@ -81,8 +84,10 @@ def paint(self, painter, option, widget):
         painter.setPen(pen)
         painter.drawPolygon(poly)
 
-        edge_rect = QtCore.QRectF(
-            rect.x() + (size * 1.7), rect.center().y() - 15, 4, 30)
+        edge_size = 30
+        edge_rect = QtCore.QRectF(rect.x() + (size * 1.6),
+                                  rect.center().y() - (edge_size / 2),
+                                  4, edge_size)
         painter.drawRect(edge_rect)
 
         painter.restore()
