add some nodes

2023-12-11 00:26:15 +09:00
parent f3f34ad4d0
commit c270d1f4b2
3 changed files with 229 additions and 71 deletions
@@ -0,0 +1,17 @@
+
+def get_node_names_mappings(classes):
+    node_names = {}
+    node_classes = {}
+    for cls in classes:
+        # check if "custom_name" attribute is set
+        if hasattr(cls, "custom_name"):
+            node_names[cls.__name__] = cls.custom_name
+            node_classes[cls.__name__] = cls
+    return node_names, node_classes
+
+def node_wrapper(container):
+    def wrap_class(cls):
+        container.append(cls)
+        return cls
+    return wrap_class
+
@@ -2,6 +2,13 @@
 Implements logic gate nodes
 """
 import re
+from .autonode import node_wrapper, get_node_names_mappings
+
+
+classes = []
+node = node_wrapper(classes)
+
+@node
 class LogicGateCompareFloat:
    """
    Returns 1 if input1 > input2, 0 otherwise
@@ -17,9 +24,167 @@ class LogicGateCompareFloat:
    }
    FUNCTION = "compareFloat"
    CATEGORY = "Logic Gates"
+    custom_name = "ABiggerThanB(Float)"
    def compareFloat(self, input1, input2):
        return 1.0 if input1 > input2 else 0.0
-    
+@node
+class LogicGateInvertBasic:
+    """
+    Inverts 1 to 0 and 0 to 1
+    """
+    RETURN_TYPES = ("INT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("INT", {"default": 0}),
+        }
+    }
+    FUNCTION = "invert"
+    CATEGORY = "Logic Gates"
+    custom_name = "Invert Basic"
+    def invert(self, input1):
+        return 1 if input1 == 0 else 0
+@node
+class LogicGateInvertValueInt:
+    """
+    Inverts x -> -x
+    """
+    RETURN_TYPES = ("INT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("INT", {"default": 0}),
+        }
+    }
+    FUNCTION = "invertValue"
+    CATEGORY = "Logic Gates"
+    custom_name = "Invert Value Int"
+    def invertValue(self, input1):
+        return -input1
+@node
+class LogicGateInvertValueFloat:
+    """
+    Inverts x -> -x
+    """
+    RETURN_TYPES = ("FLOAT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("FLOAT", {"default": 0.0}),
+        }
+    }
+    FUNCTION = "invertValue"
+    CATEGORY = "Logic Gates"
+    custom_name = "Invert Value Float"
+    def invertValue(self, input1):
+        return -input1
+@node
+class LogicGateBitwiseShift:
+    """
+    Shifts input1 by input2 bits
+    Only works on integers
+    Negative input2 shifts right, positive input2 shifts left
+    """
+    RETURN_TYPES = ("INT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("INT", {"default": 0}),
+            "input2": ("INT", {"default": 0}),
+        }
+    }
+    FUNCTION = "bitwiseShift"
+    CATEGORY = "Logic Gates"
+    custom_name = "Bitwise Shift"
+    def bitwiseShift(self, input1, input2):
+        # validate input2
+        if abs(input2) > 32:
+            raise ValueError("input2 must be between -32 and 32")
+        return input1 << input2
+@node
+class LogicGateBitwiseAnd:
+    """
+    Bitwise AND of input1 and input2
+    Only works on integers
+    """
+    RETURN_TYPES = ("INT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("INT", {"default": 0}),
+            "input2": ("INT", {"default": 0}),
+        }
+    }
+    FUNCTION = "bitwiseAnd"
+    CATEGORY = "Logic Gates"
+    custom_name = "Bitwise And"
+    def bitwiseAnd(self, input1, input2):
+        return input1 & input2
+@node
+class LogicGateBitwiseOr:
+    """
+    Bitwise OR of input1 and input2
+    Only works on integers
+    """
+    RETURN_TYPES = ("INT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("INT", {"default": 0}),
+            "input2": ("INT", {"default": 0}),
+        }
+    }
+    FUNCTION = "bitwiseOr"
+    CATEGORY = "Logic Gates"
+    custom_name = "Bitwise Or"
+    def bitwiseOr(self, input1, input2):
+        return input1 | input2
+@node
+class LogicGateBitwiseXor:
+    """
+    Bitwise XOR of input1 and input2
+    Only works on integers
+    """
+    RETURN_TYPES = ("INT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("INT", {"default": 0}),
+            "input2": ("INT", {"default": 0}),
+        }
+    }
+    FUNCTION = "bitwiseXor"
+    CATEGORY = "Logic Gates"
+    custom_name = "Bitwise Xor"
+    def bitwiseXor(self, input1, input2):
+        return input1 ^ input2
+@node
+class LogicGateBitwiseNot:
+    """
+    Bitwise NOT of input1
+    Only works on integers
+    """
+    RETURN_TYPES = ("INT",)
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+        "required": {
+            "input1": ("INT", {"default": 0}),
+        }
+    }
+    FUNCTION = "bitwiseNot"
+    CATEGORY = "Logic Gates"
+    custom_name = "Bitwise Not"
+    def bitwiseNot(self, input1):
+        return ~input1
+@node
 class LogicGateCompareInt:
    """
    Returns 1 if input1 > input2, 0 otherwise
@@ -35,9 +200,10 @@ class LogicGateCompareInt:
    }
    FUNCTION = "compareInt"
    CATEGORY = "Logic Gates"
+    custom_name = "ABiggerThanB(Int)"
    def compareInt(self, input1, input2):
        return 1 if input1 > input2 else 0
-
+@node
 class LogicGateCompareString:
    """
    Returns if given regex (1) is found in given string (2)
@@ -53,9 +219,10 @@ class LogicGateCompareString:
    }
    FUNCTION = "compareString"
    CATEGORY = "Logic Gates"
+    custom_name = "AContainsB(String)"
    def compareString(self, regex, input2):
        return 1 if re.search(regex, input2) else 0
-
+@node
 class LogicGateEitherFloat:
    """
    Returns input1 if condition is true, input2 otherwise
@@ -72,9 +239,10 @@ class LogicGateEitherFloat:
    }
    FUNCTION = "either"
    CATEGORY = "Logic Gates"
+    custom_name = "ConditionAorB(Float)"
    def either(self, condition, input1, input2):
        return input1 if condition else input2
-
+@node
 class LogicGateEitherInt:
    """
    Returns input1 if condition is true, input2 otherwise
@@ -91,9 +259,10 @@ class LogicGateEitherInt:
    }
    FUNCTION = "either"
    CATEGORY = "Logic Gates"
+    custom_name = "ConditionAorB(Int)"
    def either(self, condition, input1, input2):
        return input1 if condition else input2
-
+@node
 class StaticNumberInt:
    """
    Returns a static number
@@ -108,9 +277,10 @@ class StaticNumberInt:
    }
    FUNCTION = "staticNumber"
    CATEGORY = "Logic Gates"
+    custom_name = "Static Number Int"
    def staticNumber(self, number):
        return number
-
+@node
 class StaticNumberFloat:
    """
    Returns a static number
@@ -127,7 +297,7 @@ class StaticNumberFloat:
    CATEGORY = "Logic Gates"
    def staticNumber(self, number):
        return number
-
+@node
 class StaticString:
    """
    Returns a static string
@@ -142,9 +312,10 @@ class StaticString:
    }
    FUNCTION = "staticString"
    CATEGORY = "Logic Gates"
+    custom_name = "Static String"
    def staticString(self, string):
        return string
-
+@node
 class LogicGateAndInt:
    """
    Returns 1 if all inputs are True, 0 otherwise
@@ -160,9 +331,10 @@ class LogicGateAndInt:
    }
    FUNCTION = "and_"
    CATEGORY = "Logic Gates"
+    custom_name = "And Int"
    def and_(self, input1, input2):
        return 1 if input1 and input2 else 0
-
+@node
 class LogicGateAndFloat:
    """
    Returns 1 if all inputs are True, 0 otherwise
@@ -178,8 +350,10 @@ class LogicGateAndFloat:
    }
    FUNCTION = "and_"
    CATEGORY = "Logic Gates"
+    custom_name = "And Float"
    def and_(self, input1, input2):
        return 1 if input1 and input2 else 0
+@node
 class LogicGateOrInt:
    """
    Returns 1 if any input is True, 0 otherwise
@@ -195,8 +369,10 @@ class LogicGateOrInt:
    }
    FUNCTION = "or_"
    CATEGORY = "Logic Gates"
+    custom_name = "Or Int"
    def or_(self, input1, input2):
        return 1 if input1 or input2 else 0
+@node
 class LogicGateOrFloat:
    """
    Returns 1 if any input is True, 0 otherwise
@@ -212,9 +388,10 @@ class LogicGateOrFloat:
    }
    FUNCTION = "or_"
    CATEGORY = "Logic Gates"
+    custom_name = "Or Float"
    def or_(self, input1, input2):
        return 1 if input1 or input2 else 0    
-
+@node
 class LogicGateEitherString:
    """
    Returns input1 if condition is true, input2 otherwise
@@ -231,9 +408,10 @@ class LogicGateEitherString:
    }
    FUNCTION = "either"
    CATEGORY = "Logic Gates"
+    custom_name = "Either String"
    def either(self, condition, input1, input2):
        return input1 if condition else input2
-
+@node
 class AddInt:
    """
    Returns the sum of the inputs
@@ -249,9 +427,10 @@ class AddInt:
    }
    FUNCTION = "add"
    CATEGORY = "Logic Gates"
+    custom_name = "Add Int"
    def add(self, input1, input2):
        return input1 + input2
-
+@node
 class AddFloat:
    """
    Returns the sum of the inputs
@@ -267,9 +446,10 @@ class AddFloat:
    }
    FUNCTION = "add"
    CATEGORY = "Logic Gates"
+    custom_name = "Add Float"
    def add(self, input1, input2):
        return input1 + input2
-
+@node
 class MergeString:
    """
    Returns the concatenation of the inputs
@@ -285,43 +465,8 @@ class MergeString:
    }
    FUNCTION = "merge"
    CATEGORY = "Logic Gates"
+    custom_name = "Merge String"
    def merge(self, input1, input2):
        return input1 + input2

-CLASS_MAPPINGS = {
-    "LogicGateCompareFloat": LogicGateCompareFloat,
-    "LogicGateCompareInt": LogicGateCompareInt,
-    "LogicGateCompareString": LogicGateCompareString,
-    "LogicGateEitherFloat": LogicGateEitherFloat,
-    "LogicGateEitherInt": LogicGateEitherInt,
-    "StaticNumberInt": StaticNumberInt,
-    "StaticNumberFloat": StaticNumberFloat,
-    "StaticString": StaticString,
-    "LogicGateAndInt": LogicGateAndInt,
-    "LogicGateAndFloat": LogicGateAndFloat,
-    "LogicGateOrInt": LogicGateOrInt,
-    "LogicGateOrFloat": LogicGateOrFloat,
-    "LogicGateEitherString": LogicGateEitherString,
-    "AddInt": AddInt,
-    "AddFloat": AddFloat,
-    "MergeString": MergeString,
-}
-
-CLASS_NAMES = {
-    "LogicGateCompareFloat": "ABiggerThanB(Float)",
-    "LogicGateCompareInt": "ABiggerThanB(Int)",
-    "LogicGateCompareString": "AContainsB(String)",
-    "LogicGateEitherFloat": "ConditionAorB(Float)",
-    "LogicGateEitherInt": "ConditionAorB(Int)",
-    "StaticNumberInt": "Static Number Int",
-    "StaticNumberFloat": "Static Number Float",
-    "StaticString": "Static String",
-    "LogicGateAndInt": "And(Int)",
-    "LogicGateAndFloat": "And(Float)",
-    "LogicGateOrInt": "Or(Int)",
-    "LogicGateOrFloat": "Or(Float)",
-    "LogicGateEitherString": "Either String",
-    "AddInt": "Add Int",
-    "AddFloat": "Add Float",
-    "MergeString": "Merge String",
-}
+CLASS_MAPPINGS, CLASS_NAMES = get_node_names_mappings(classes)
@@ -1,5 +1,10 @@
 import random
+from .autonode import node_wrapper, get_node_names_mappings

+classes = []
+node = node_wrapper(classes)
+
+@node
 class UniformRandomFloat:
    """
    Selects a random float from min to max
@@ -29,7 +34,8 @@ class UniformRandomFloat:
    RETURN_TYPES = ("FLOAT",)
    FUNCTION = "generate"
    CATEGORY = "Logic Gates"
-
+    custom_name = "Uniform Random Float"
+@node
 class UniformRandomInt:
    """
    Selects a random int from min to max
@@ -56,7 +62,8 @@ class UniformRandomInt:
    RETURN_TYPES = ("INT",)
    FUNCTION = "generate"
    CATEGORY = "Logic Gates"
-
+    custom_name = "Uniform Random Int"
+@node
 class UniformRandomChoice:
    """
    Parses input string with separator '$' and returns a random choice
@@ -81,7 +88,8 @@ class UniformRandomChoice:
    RETURN_TYPES = ("STRING",)
    FUNCTION = "generate"
    CATEGORY = "Logic Gates"
-
+    custom_name = "Uniform Random Choice"
+@node
 class ManualChoiceString:
    """
    Parses input string with separator '$' and returns a random choice
@@ -106,7 +114,8 @@ class ManualChoiceString:
    RETURN_TYPES = ("STRING",)
    FUNCTION = "generate"
    CATEGORY = "Logic Gates"
-
+    custom_name = "Manual Choice String"
+@node
 class ManualChoiceInt:
    """
    Parses input string with separator '$' and returns a random choice
@@ -132,7 +141,8 @@ class ManualChoiceInt:
    RETURN_TYPES = ("INT",)
    FUNCTION = "generate"
    CATEGORY = "Logic Gates"
-
+    custom_name = "Manual Choice Int"
+@node
 class ManualChoiceFloat:
    """
    Parses input string with separator '$' and returns a random choice
@@ -158,20 +168,6 @@ class ManualChoiceFloat:
    RETURN_TYPES = ("FLOAT",)
    FUNCTION = "generate"
    CATEGORY = "Logic Gates"
+    custom_name = "Manual Choice Float"

-CLASS_MAPPINGS = {
-    "UniformRandomFloat": UniformRandomFloat,
-    "UniformRandomInt": UniformRandomInt,
-    "UniformRandomChoice": UniformRandomChoice,
-    "ManualChoiceString": ManualChoiceString,
-    "ManualChoiceInt": ManualChoiceInt,
-    "ManualChoiceFloat": ManualChoiceFloat,
-}
-CLASS_NAMES = {
-    "UniformRandomFloat": "Uniform Random Float",
-    "UniformRandomInt": "Uniform Random Int",
-    "UniformRandomChoice": "Uniform Random Choice",
-    "ManualChoiceString": "Manual Choice String",
-    "ManualChoiceInt": "Manual Choice Int",
-    "ManualChoiceFloat": "Manual Choice Float",
-}
+CLASS_MAPPINGS, CLASS_NAMES = get_node_names_mappings(classes)