ComfyUI/comfy_extras/nodes/nodes_regexp.py
2025-02-18 20:23:09 -08:00

108 lines
2.9 KiB
Python

import re
from comfy.node_helpers import export_custom_nodes
from comfy.nodes.package_typing import CustomNode, InputTypes
MATCH_TYPE_NAME = "MATCH"
class RegexFlags(CustomNode):
@classmethod
def INPUT_TYPES(cls) -> InputTypes:
flags_ = {flag_name: ("BOOLEAN", {}) for flag_name in ["ASCII", "IGNORECASE", "LOCALE", "MULTILINE", "DOTALL", "VERBOSE", "UNICODE"]}
return {
"required": flags_,
}
CATEGORY = "regular_expressions"
FUNCTION = "execute"
RETURN_TYPES = ("INT",)
def execute(self, **kwargs) -> tuple[int]:
has_noflag = hasattr(re.RegexFlag, "NOFLAG")
# use getattr for python 3.10 compatibility
flags = getattr(re.RegexFlag, "NOFLAG") if has_noflag else 0
for name, on in kwargs.items():
if on and hasattr(re.RegexFlag, name):
flags |= int(re.RegexFlag[name])
return int(flags),
class Regex(CustomNode):
@classmethod
def INPUT_TYPES(cls) -> InputTypes:
return {
"required": {
"pattern": ("STRING", {}),
"string": ("STRING", {}),
},
"optional": {
"flags": ("INT", {"min": 0, "default": 0})
}
}
CATEGORY = "regular_expressions"
FUNCTION = "execute"
RETURN_TYPES = (MATCH_TYPE_NAME,)
def execute(self, pattern: str = "", string: str = "", flags: int = 0) -> tuple[re.Match]:
return re.match(pattern=pattern, string=string, flags=flags),
class RegexMatchGroupByIndex(CustomNode):
@classmethod
def INPUT_TYPES(cls) -> InputTypes:
return {
"required": {
"match": (MATCH_TYPE_NAME, {}),
"index": ("INT", {}),
},
}
CATEGORY = "regular_expressions"
FUNCTION = "execute"
RETURN_TYPES = (MATCH_TYPE_NAME,)
def execute(self, match: re.Match, index: int = 0) -> tuple[str]:
return match.group(index),
class RegexMatchGroupByName(CustomNode):
@classmethod
def INPUT_TYPES(cls) -> InputTypes:
return {
"required": {
"match": (MATCH_TYPE_NAME, {}),
"name": ("STRING", {}),
},
}
CATEGORY = "regular_expressions"
FUNCTION = "execute"
RETURN_TYPES = (MATCH_TYPE_NAME,)
def execute(self, match: re.Match, name: str = "") -> tuple[str]:
return match.group(name),
class RegexMatchExpand(CustomNode):
@classmethod
def INPUT_TYPES(cls) -> InputTypes:
return {
"required": {
"match": (MATCH_TYPE_NAME, {}),
"template": ("STRING", {}),
},
}
CATEGORY = "regular_expressions"
FUNCTION = "execute"
RETURN_TYPES = ("STRING",)
def execute(self, match: re.Match, template: str = "") -> tuple[str]:
return match.expand(template),
export_custom_nodes()