Procházet zdrojové kódy

Add Comparison type socket to Compare node

Joseph Brandenburg před 6 měsíci
rodič
revize
625488d9ad
5 změnil soubory, kde provedl 76 přidání a 27 odebrání
  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'),
                ("LinkCopyScale",             'INPUT',  "BooleanSocket", "Additive",     3,      False,    False),
                ("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_types = ["NodeGroupInput", "NodeGroupOutput", "SchemaIncomingConnection",

+ 21 - 21
misc_nodes.py

@@ -1,6 +1,5 @@
 from .node_container_common import *
 from .base_definitions import MantisNode, NodeSocket
-from .base_definitions import MantisSocketTemplate as SockTemplate
 from .xForm_containers import xFormArmature, xFormBone
 from .misc_nodes_socket_templates import *
 from math import pi, tau
@@ -1534,13 +1533,6 @@ class UtilityMatrixTransform(MantisNode):
         self.prepared = 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):
     def __init__(self, signature, base_tree):
         super().__init__(signature, base_tree, MatrixInvertSockets)
@@ -1727,27 +1719,35 @@ class UtilityPrint(MantisNode):
         #     prRed("No input to print.")
         self.executed = True
 
-
 class UtilityCompare(MantisNode):
     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.node_type = "UTILITY"
 
     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
 
-
 class UtilityChoose(MantisNode):
     def __init__(self, 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 dataclasses import replace
 
 GetCurvePointSockets=[
     CurveTemplate := SockTemplate(name="Curve", bl_idname='EnumCurveSocket', 
@@ -23,3 +24,21 @@ GetNearestFactorOnCurveSockets=[
     OutputFactorTemplate := SockTemplate(name="Factor",
         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.initialized = True
 
-from .misc_nodes import MatrixInvertSockets
 class UtilityMatrixInvert(Node, MantisUINode):
     """Inverts an invertable matrix, otherwise throws an error."""
     bl_idname = "UtilityMatrixInvert"
@@ -1089,9 +1088,7 @@ class UtilityCompare(Node, MantisUINode):
     mantis_node_class_name=bl_idname
     
     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
     
     def update(self):

+ 32 - 1
socket_definitions.py

@@ -194,7 +194,8 @@ def TellClasses() -> List[MantisSocket]:
              MathFloatOperation,
              MathVectorOperation,
              MatrixTransformOperation,
-
+             #conditions
+             EnumCompareOperation,
              # Schema
              WildcardSocket,
             #  xFormArraySocket,
@@ -2402,6 +2403,36 @@ class MathIntOperation(MantisSocket):
         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):
     """Some kind of node socket lol I donno"""
     bl_idname = 'WildcardSocket'