convert String nodes to V3 schema (#9370)

comfy_extras/nodes_string.py

@@ -1,77 +1,91 @@
 import re
+from typing_extensions import override
 
-from comfy.comfy_types.node_typing import IO
+from comfy_api.latest import ComfyExtension, io
 
-class StringConcatenate():
+
+class StringConcatenate(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string_a": (IO.STRING, {"multiline": True}),
-                "string_b": (IO.STRING, {"multiline": True}),
-                "delimiter": (IO.STRING, {"multiline": False, "default": ""})
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StringConcatenate",
+            display_name="Concatenate",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string_a", multiline=True),
+                io.String.Input("string_b", multiline=True),
+                io.String.Input("delimiter", multiline=False, default=""),
+            ],
+            outputs=[
+                io.String.Output(),
+            ]
+        )
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string_a, string_b, delimiter, **kwargs):
-        return delimiter.join((string_a, string_b)),
-
-class StringSubstring():
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "start": (IO.INT, {}),
-                "end": (IO.INT, {}),
-            }
-        }
+    def execute(cls, string_a, string_b, delimiter):
+        return io.NodeOutput(delimiter.join((string_a, string_b)))
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
 
-    def execute(self, string, start, end, **kwargs):
-        return string[start:end],
-
-class StringLength():
+class StringSubstring(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True})
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StringSubstring",
+            display_name="Substring",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.Int.Input("start"),
+                io.Int.Input("end"),
+            ],
+            outputs=[
+                io.String.Output(),
+            ]
+        )
 
-    RETURN_TYPES = (IO.INT,)
-    RETURN_NAMES = ("length",)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string, **kwargs):
-        length = len(string)
-
-        return length,
-
-class CaseConverter():
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "mode": (IO.COMBO, {"options": ["UPPERCASE", "lowercase", "Capitalize", "Title Case"]})
-            }
-        }
+    def execute(cls, string, start, end):
+        return io.NodeOutput(string[start:end])
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
 
-    def execute(self, string, mode, **kwargs):
+class StringLength(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StringLength",
+            display_name="Length",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+            ],
+            outputs=[
+                io.Int.Output(display_name="length"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, string):
+        return io.NodeOutput(len(string))
+
+
+class CaseConverter(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="CaseConverter",
+            display_name="Case Converter",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.Combo.Input("mode", options=["UPPERCASE", "lowercase", "Capitalize", "Title Case"]),
+            ],
+            outputs=[
+                io.String.Output(),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, string, mode):
         if mode == "UPPERCASE":
             result = string.upper()
         elif mode == "lowercase":
@@ -83,24 +97,27 @@ class CaseConverter():
         else:
             result = string
 
-        return result,
+        return io.NodeOutput(result)
 
 
-class StringTrim():
+class StringTrim(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "mode": (IO.COMBO, {"options": ["Both", "Left", "Right"]})
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StringTrim",
+            display_name="Trim",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.Combo.Input("mode", options=["Both", "Left", "Right"]),
+            ],
+            outputs=[
+                io.String.Output(),
+            ]
+        )
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string, mode, **kwargs):
+    @classmethod
+    def execute(cls, string, mode):
         if mode == "Both":
             result = string.strip()
         elif mode == "Left":
@@ -110,70 +127,78 @@ class StringTrim():
         else:
             result = string
 
-        return result,
+        return io.NodeOutput(result)
 
-class StringReplace():
+
+class StringReplace(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "find": (IO.STRING, {"multiline": True}),
-                "replace": (IO.STRING, {"multiline": True})
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StringReplace",
+            display_name="Replace",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.String.Input("find", multiline=True),
+                io.String.Input("replace", multiline=True),
+            ],
+            outputs=[
+                io.String.Output(),
+            ]
+        )
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string, find, replace, **kwargs):
-        result = string.replace(find, replace)
-        return result,
-
-
-class StringContains():
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "substring": (IO.STRING, {"multiline": True}),
-                "case_sensitive": (IO.BOOLEAN, {"default": True})
-            }
-        }
+    def execute(cls, string, find, replace):
+        return io.NodeOutput(string.replace(find, replace))
 
-    RETURN_TYPES = (IO.BOOLEAN,)
-    RETURN_NAMES = ("contains",)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
 
-    def execute(self, string, substring, case_sensitive, **kwargs):
+class StringContains(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StringContains",
+            display_name="Contains",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.String.Input("substring", multiline=True),
+                io.Boolean.Input("case_sensitive", default=True),
+            ],
+            outputs=[
+                io.Boolean.Output(display_name="contains"),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, string, substring, case_sensitive):
         if case_sensitive:
             contains = substring in string
         else:
             contains = substring.lower() in string.lower()
 
-        return contains,
+        return io.NodeOutput(contains)
 
 
-class StringCompare():
+class StringCompare(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string_a": (IO.STRING, {"multiline": True}),
-                "string_b": (IO.STRING, {"multiline": True}),
-                "mode": (IO.COMBO, {"options": ["Starts With", "Ends With", "Equal"]}),
-                "case_sensitive": (IO.BOOLEAN, {"default": True})
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="StringCompare",
+            display_name="Compare",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string_a", multiline=True),
+                io.String.Input("string_b", multiline=True),
+                io.Combo.Input("mode", options=["Starts With", "Ends With", "Equal"]),
+                io.Boolean.Input("case_sensitive", default=True),
+            ],
+            outputs=[
+                io.Boolean.Output(),
+            ]
+        )
 
-    RETURN_TYPES = (IO.BOOLEAN,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string_a, string_b, mode, case_sensitive, **kwargs):
+    @classmethod
+    def execute(cls, string_a, string_b, mode, case_sensitive):
         if case_sensitive:
             a = string_a
             b = string_b
@@ -182,31 +207,34 @@ class StringCompare():
             b = string_b.lower()
 
         if mode == "Equal":
-            return a == b,
+            return io.NodeOutput(a == b)
         elif mode == "Starts With":
-            return a.startswith(b),
+            return io.NodeOutput(a.startswith(b))
         elif mode == "Ends With":
-            return a.endswith(b),
+            return io.NodeOutput(a.endswith(b))
 
-class RegexMatch():
+
+class RegexMatch(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "regex_pattern": (IO.STRING, {"multiline": True}),
-                "case_insensitive": (IO.BOOLEAN, {"default": True}),
-                "multiline": (IO.BOOLEAN, {"default": False}),
-                "dotall": (IO.BOOLEAN, {"default": False})
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="RegexMatch",
+            display_name="Regex Match",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.String.Input("regex_pattern", multiline=True),
+                io.Boolean.Input("case_insensitive", default=True),
+                io.Boolean.Input("multiline", default=False),
+                io.Boolean.Input("dotall", default=False),
+            ],
+            outputs=[
+                io.Boolean.Output(display_name="matches"),
+            ]
+        )
 
-    RETURN_TYPES = (IO.BOOLEAN,)
-    RETURN_NAMES = ("matches",)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string, regex_pattern, case_insensitive, multiline, dotall, **kwargs):
+    @classmethod
+    def execute(cls, string, regex_pattern, case_insensitive, multiline, dotall):
         flags = 0
 
         if case_insensitive:
@@ -223,29 +251,32 @@ class RegexMatch():
         except re.error:
             result = False
 
-        return result,
+        return io.NodeOutput(result)
 
 
-class RegexExtract():
+class RegexExtract(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "regex_pattern": (IO.STRING, {"multiline": True}),
-                "mode": (IO.COMBO, {"options": ["First Match", "All Matches", "First Group", "All Groups"]}),
-                "case_insensitive": (IO.BOOLEAN, {"default": True}),
-                "multiline": (IO.BOOLEAN, {"default": False}),
-                "dotall": (IO.BOOLEAN, {"default": False}),
-                "group_index": (IO.INT, {"default": 1, "min": 0, "max": 100})
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="RegexExtract",
+            display_name="Regex Extract",
+            category="utils/string",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.String.Input("regex_pattern", multiline=True),
+                io.Combo.Input("mode", options=["First Match", "All Matches", "First Group", "All Groups"]),
+                io.Boolean.Input("case_insensitive", default=True),
+                io.Boolean.Input("multiline", default=False),
+                io.Boolean.Input("dotall", default=False),
+                io.Int.Input("group_index", default=1, min=0, max=100),
+            ],
+            outputs=[
+                io.String.Output(),
+            ]
+        )
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string, regex_pattern, mode, case_insensitive, multiline, dotall, group_index, **kwargs):
+    @classmethod
+    def execute(cls, string, regex_pattern, mode, case_insensitive, multiline, dotall, group_index):
         join_delimiter = "\n"
 
         flags = 0
@@ -294,32 +325,33 @@ class RegexExtract():
         except re.error:
             result = ""
 
-        return result,
+        return io.NodeOutput(result)
 
 
-class RegexReplace():
-    DESCRIPTION = "Find and replace text using regex patterns."
+class RegexReplace(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "string": (IO.STRING, {"multiline": True}),
-                "regex_pattern": (IO.STRING, {"multiline": True}),
-                "replace": (IO.STRING, {"multiline": True}),
-            },
-            "optional": {
-                "case_insensitive": (IO.BOOLEAN, {"default": True}),
-                "multiline": (IO.BOOLEAN, {"default": False}),
-                "dotall": (IO.BOOLEAN, {"default": False, "tooltip": "When enabled, the dot (.) character will match any character including newline characters. When disabled, dots won't match newlines."}),
-                "count": (IO.INT, {"default": 0, "min": 0, "max": 100, "tooltip": "Maximum number of replacements to make. Set to 0 to replace all occurrences (default). Set to 1 to replace only the first match, 2 for the first two matches, etc."}),
-            }
-        }
+    def define_schema(cls):
+        return io.Schema(
+            node_id="RegexReplace",
+            display_name="Regex Replace",
+            category="utils/string",
+            description="Find and replace text using regex patterns.",
+            inputs=[
+                io.String.Input("string", multiline=True),
+                io.String.Input("regex_pattern", multiline=True),
+                io.String.Input("replace", multiline=True),
+                io.Boolean.Input("case_insensitive", default=True, optional=True),
+                io.Boolean.Input("multiline", default=False, optional=True),
+                io.Boolean.Input("dotall", default=False, optional=True, tooltip="When enabled, the dot (.) character will match any character including newline characters. When disabled, dots won't match newlines."),
+                io.Int.Input("count", default=0, min=0, max=100, optional=True, tooltip="Maximum number of replacements to make. Set to 0 to replace all occurrences (default). Set to 1 to replace only the first match, 2 for the first two matches, etc."),
+            ],
+            outputs=[
+                io.String.Output(),
+            ]
+        )
 
-    RETURN_TYPES = (IO.STRING,)
-    FUNCTION = "execute"
-    CATEGORY = "utils/string"
-
-    def execute(self, string, regex_pattern, replace, case_insensitive=True, multiline=False, dotall=False, count=0, **kwargs):
+    @classmethod
+    def execute(cls, string, regex_pattern, replace, case_insensitive=True, multiline=False, dotall=False, count=0):
         flags = 0
 
         if case_insensitive:
@@ -329,32 +361,25 @@ class RegexReplace():
         if dotall:
             flags |= re.DOTALL
         result = re.sub(regex_pattern, replace, string, count=count, flags=flags)
-        return result,
+        return io.NodeOutput(result)
 
-NODE_CLASS_MAPPINGS = {
-    "StringConcatenate": StringConcatenate,
-    "StringSubstring": StringSubstring,
-    "StringLength": StringLength,
-    "CaseConverter": CaseConverter,
-    "StringTrim": StringTrim,
-    "StringReplace": StringReplace,
-    "StringContains": StringContains,
-    "StringCompare": StringCompare,
-    "RegexMatch": RegexMatch,
-    "RegexExtract": RegexExtract,
-    "RegexReplace": RegexReplace,
-}
 
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "StringConcatenate": "Concatenate",
-    "StringSubstring": "Substring",
-    "StringLength": "Length",
-    "CaseConverter": "Case Converter",
-    "StringTrim": "Trim",
-    "StringReplace": "Replace",
-    "StringContains": "Contains",
-    "StringCompare": "Compare",
-    "RegexMatch": "Regex Match",
-    "RegexExtract": "Regex Extract",
-    "RegexReplace": "Regex Replace",
-}
+class StringExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            StringConcatenate,
+            StringSubstring,
+            StringLength,
+            CaseConverter,
+            StringTrim,
+            StringReplace,
+            StringContains,
+            StringCompare,
+            RegexMatch,
+            RegexExtract,
+            RegexReplace,
+        ]
+
+async def comfy_entrypoint() -> StringExtension:
+    return StringExtension()