Browse Source

Add Comparison type socket to Compare node

Joseph Brandenburg 6 months ago
parent
commit
625488d9ad
5 changed files with 76 additions and 27 deletions
  1. 3 1
      base_definitions.py
  2. 21 21
      misc_nodes.py
  3. 19 0
      misc_nodes_socket_templates.py
  4. 1 4
      misc_nodes_ui.py
  5. 32 1
      socket_definitions.py

+ 3 - 1
base_definitions.py

@@ -582,7 +582,9 @@ SOCKETS_ADDED=[("DeformerMorphTargetDeform", 'INPUT', 'BooleanSocket', "Use Shap
                ("UtilityFCurve",             'INPUT',  "eFCrvExtrapolationMode", "Extrapolation Mode", 0, False, 'CONSTANT'),
                ("UtilityFCurve",             'INPUT',  "eFCrvExtrapolationMode", "Extrapolation Mode", 0, False, 'CONSTANT'),
                ("LinkCopyScale",             'INPUT',  "BooleanSocket", "Additive",     3,      False,    False),
                ("LinkCopyScale",             'INPUT',  "BooleanSocket", "Additive",     3,      False,    False),
                ("DeformerHook",              'INPUT',  "FloatFactorSocket", "Influence", 3,     False,    1.0),
                ("DeformerHook",              'INPUT',  "FloatFactorSocket", "Influence", 3,     False,    1.0),
-               ("DeformerHook",              'INPUT',  "BooleanSocket", "Auto-Bezier",   4,     False,    True),]
+               ("DeformerHook",              'INPUT',  "BooleanSocket", "Auto-Bezier",   4,     False,    True),
+               ("UtilityCompare",            'INPUT',  "EnumCompareOperation", "Comparison", 0,  False,   'EQUAL'),
+               ]
 
 
 # replace names with bl_idnames for reading the tree and solving schemas.
 # replace names with bl_idnames for reading the tree and solving schemas.
 replace_types = ["NodeGroupInput", "NodeGroupOutput", "SchemaIncomingConnection",
 replace_types = ["NodeGroupInput", "NodeGroupOutput", "SchemaIncomingConnection",

+ 21 - 21
misc_nodes.py

@@ -1,6 +1,5 @@
 from .node_container_common import *
 from .node_container_common import *
 from .base_definitions import MantisNode, NodeSocket
 from .base_definitions import MantisNode, NodeSocket
-from .base_definitions import MantisSocketTemplate as SockTemplate
 from .xForm_containers import xFormArmature, xFormBone
 from .xForm_containers import xFormArmature, xFormBone
 from .misc_nodes_socket_templates import *
 from .misc_nodes_socket_templates import *
 from math import pi, tau
 from math import pi, tau
@@ -1534,13 +1533,6 @@ class UtilityMatrixTransform(MantisNode):
         self.prepared = True
         self.prepared = True
         self.executed = True
         self.executed = True
 
 
-MatrixInvertSockets=[
-    Matrix1Template := SockTemplate(
-    name="Matrix 1", is_input=True,  bl_idname='MatrixSocket', ),
-    MatrixOutTemplate := SockTemplate(
-    name="Matrix", is_input=False,  bl_idname='MatrixSocket', ),
-]
-
 class UtilityMatrixInvert(MantisNode):
 class UtilityMatrixInvert(MantisNode):
     def __init__(self, signature, base_tree):
     def __init__(self, signature, base_tree):
         super().__init__(signature, base_tree, MatrixInvertSockets)
         super().__init__(signature, base_tree, MatrixInvertSockets)
@@ -1727,27 +1719,35 @@ class UtilityPrint(MantisNode):
         #     prRed("No input to print.")
         #     prRed("No input to print.")
         self.executed = True
         self.executed = True
 
 
-
 class UtilityCompare(MantisNode):
 class UtilityCompare(MantisNode):
     def __init__(self, signature, base_tree):
     def __init__(self, signature, base_tree):
-        super().__init__(signature, base_tree)
-        inputs = [
-          "A"           ,
-          "B"           ,
-        ]
-        outputs = [
-          "Result"      ,
-        ]
-        self.inputs.init_sockets(inputs)
-        self.outputs.init_sockets(outputs)
+        super().__init__(signature, base_tree, CompareSockets)
         self.init_parameters()
         self.init_parameters()
         self.node_type = "UTILITY"
         self.node_type = "UTILITY"
 
 
     def bPrepare(self, bContext = None,):
     def bPrepare(self, bContext = None,):
-        self.parameters["Result"] = self.evaluate_input("A") == self.evaluate_input("B")
+        operation=self.evaluate_input("Comparison")
+        a = self.evaluate_input("A")
+        b = self.evaluate_input("B")
+        if isinstance(a, str) and isinstance(b, str) and \
+              operation not in ['EQUAL', 'NOT_EQUAL']:
+                   raise GraphError("Strings do not have numerical value to"
+                                    " compute greater than or less than.")
+        match operation:
+            case "EQUAL":
+                self.parameters["Result"] = a == b
+            case "NOT_EQUAL":
+                self.parameters["Result"] = a != b
+            case "GREATER_THAN":
+                self.parameters["Result"] = a > b
+            case "GREATER_THAN_EQUAL":
+                self.parameters["Result"] = a >= b
+            case "LESS_THAN":
+                self.parameters["Result"] = a < b
+            case "LESS_THAN_EQUAL":
+                self.parameters["Result"] = a <= b
         self.prepared = True; self.executed = True
         self.prepared = True; self.executed = True
 
 
-
 class UtilityChoose(MantisNode):
 class UtilityChoose(MantisNode):
     def __init__(self, signature, base_tree):
     def __init__(self, signature, base_tree):
         super().__init__(signature, base_tree)
         super().__init__(signature, base_tree)

+ 19 - 0
misc_nodes_socket_templates.py

@@ -1,4 +1,5 @@
 from .base_definitions import MantisSocketTemplate as SockTemplate
 from .base_definitions import MantisSocketTemplate as SockTemplate
+from dataclasses import replace
 
 
 GetCurvePointSockets=[
 GetCurvePointSockets=[
     CurveTemplate := SockTemplate(name="Curve", bl_idname='EnumCurveSocket', 
     CurveTemplate := SockTemplate(name="Curve", bl_idname='EnumCurveSocket', 
@@ -23,3 +24,21 @@ GetNearestFactorOnCurveSockets=[
     OutputFactorTemplate := SockTemplate(name="Factor",
     OutputFactorTemplate := SockTemplate(name="Factor",
         bl_idname='FloatSocket', is_input=False,),
         bl_idname='FloatSocket', is_input=False,),
 ]
 ]
+
+MatrixInvertSockets=[
+    Matrix1Template := SockTemplate(
+    name="Matrix 1", is_input=True,  bl_idname='MatrixSocket', ),
+    MatrixOutTemplate := SockTemplate(
+    name="Matrix", is_input=False,  bl_idname='MatrixSocket', ),
+]
+
+CompareSockets = [
+    ComparisonOperation := SockTemplate( name='Comparison',
+            is_input=True, bl_idname="EnumCompareOperation",
+            default_value="EQUAL",),
+    WildcardATemplate := SockTemplate(
+        name="A", is_input=True,  bl_idname='WildcardSocket', ),
+    WildcardBTemplate := replace(WildcardATemplate, name="B"),
+    CompareOutputTemplate := SockTemplate(
+        name="Result", is_input=False, bl_idname="BooleanSocket",),
+]

+ 1 - 4
misc_nodes_ui.py

@@ -880,7 +880,6 @@ class UtilityMatrixTransform(Node, MantisUINode):
         self.outputs.new("MatrixSocket", "Out Matrix")
         self.outputs.new("MatrixSocket", "Out Matrix")
         self.initialized = True
         self.initialized = True
 
 
-from .misc_nodes import MatrixInvertSockets
 class UtilityMatrixInvert(Node, MantisUINode):
 class UtilityMatrixInvert(Node, MantisUINode):
     """Inverts an invertable matrix, otherwise throws an error."""
     """Inverts an invertable matrix, otherwise throws an error."""
     bl_idname = "UtilityMatrixInvert"
     bl_idname = "UtilityMatrixInvert"
@@ -1089,9 +1088,7 @@ class UtilityCompare(Node, MantisUINode):
     mantis_node_class_name=bl_idname
     mantis_node_class_name=bl_idname
     
     
     def init(self, context):
     def init(self, context):
-        self.inputs.new("WildcardSocket", "A")
-        self.inputs.new("WildcardSocket", "B")
-        self.outputs.new("BooleanSocket", "Result")
+        self.init_sockets(CompareSockets)
         self.initialized = True
         self.initialized = True
     
     
     def update(self):
     def update(self):

+ 32 - 1
socket_definitions.py

@@ -194,7 +194,8 @@ def TellClasses() -> List[MantisSocket]:
              MathFloatOperation,
              MathFloatOperation,
              MathVectorOperation,
              MathVectorOperation,
              MatrixTransformOperation,
              MatrixTransformOperation,
-
+             #conditions
+             EnumCompareOperation,
              # Schema
              # Schema
              WildcardSocket,
              WildcardSocket,
             #  xFormArraySocket,
             #  xFormArraySocket,
@@ -2402,6 +2403,36 @@ class MathIntOperation(MantisSocket):
         return self.color_simple
         return self.color_simple
 
 
 
 
+enumCompareOperations =  (('EQUAL', "Equal", "Equal"),
+                          ('NOT_EQUAL', "Not Equal", "Not Equal"),
+                          ('GREATER_THAN', "Greater Than", "Greater Than"),
+                          ('GREATER_THAN_EQUAL', "Greater Than or Equal", "Greater Than or Equal"),
+                          ('LESS_THAN', "Less Than", "Less Than"),
+                          ('LESS_THAN_EQUAL', "Equal or Less Than", "Equal or Less Than"),)
+
+class EnumCompareOperation(MantisSocket):
+    """Compare Operation"""
+    bl_idname = 'EnumCompareOperation'
+    bl_label = "Comparison"
+    
+    default_value :bpy.props.EnumProperty(
+        name="",
+        description="Comparison",
+        items=enumCompareOperations,
+        default='EQUAL',
+        update = update_socket,)
+    
+    color_simple = cString
+    color : bpy.props.FloatVectorProperty(default=cString, size=4)
+    input : bpy.props.BoolProperty(default =False,)
+    def draw(self, context, layout, node, text):
+        ChooseDraw(self, context, layout, node, text)
+    def draw_color(self, context, node):
+        return self.color
+    @classmethod
+    def draw_color_simple(self):
+        return self.color_simple
+
 class WildcardSocket(MantisSocket):
 class WildcardSocket(MantisSocket):
     """Some kind of node socket lol I donno"""
     """Some kind of node socket lol I donno"""
     bl_idname = 'WildcardSocket'
     bl_idname = 'WildcardSocket'