From 0d9dfd2e05b03d4bdeb8fae6b62cb8bfb81cdd27 Mon Sep 17 00:00:00 2001 From: MaraScott Date: Thu, 28 Aug 2025 12:35:38 +0200 Subject: [PATCH] Update node_formatter.py to accept curly bracket in template for non variables in dict --- node_formatter.py | 235 +++++++++++++++++++++++++++------------------- 1 file changed, 141 insertions(+), 94 deletions(-) diff --git a/node_formatter.py b/node_formatter.py index ce2210d..ef6f7d6 100644 --- a/node_formatter.py +++ b/node_formatter.py @@ -1,125 +1,172 @@ # encoding: utf-8 """ +Updated StringConstructorFormatter with template sanitization """ import typing as _t - +import re as _re from inspect import cleandoc as _cleandoc import sys as _sys from frozendict import deepfreeze as _deepfreeze - from comfy.comfy_types.node_typing import IO as _IO from . import _meta from .docstring_formatter import format_docstring as _format_docstring from .enums import DataTypes as _DataTypes from .funcs_common import _show_text_on_node, _verify_input_dict -# from .funcs_formatter import formatter as _formatter +_RECURSION_LIMIT = max(int(_sys.getrecursionlimit()), 1) -_RECURSION_LIMIT = max(int(_sys.getrecursionlimit()), 1) # You can externally monkey-patch it... but if it blows up, your fault 🤷🏻‍♂️ +def _safe_format(template: str, format_dict: _t.Dict[str, _t.Any]) -> str: + """ + Safe format that only replaces variables that exist in the dictionary. + Leaves other curly brackets untouched. + """ + if not format_dict: + return template + + # Create pattern that matches {key} only for keys that exist in format_dict + pattern = r'\{(' + '|'.join(_re.escape(key) for key in format_dict.keys()) + r')\}' + + def replace_func(match): + key = match.group(1) + return str(format_dict[key]) + + return _re.sub(pattern, replace_func, template) +def _escape_unknown_brackets(template: str, format_dict: _t.Dict[str, _t.Any]) -> str: + """ + Escape curly brackets that aren't format variables, then use normal format. + This approach preserves the original behavior for known variables. + """ + if not format_dict: + # If no format dict, escape all curly brackets + return template.replace('{', '{{').replace('}', '}}') + + # Find all {word} patterns + pattern = r'\{([^{}]*)\}' + known_keys = set(format_dict.keys()) + + def escape_func(match): + key = match.group(1) + if key in known_keys: + return match.group(0) # Keep as is for formatting + else: + return '{{' + key + '}}' # Escape unknown + + escaped_template = _re.sub(pattern, escape_func, template) + return escaped_template.format_map(format_dict) -def _recursive_format(template: str, format_dict: _t.Dict[str, _t.Any], show: bool = True, unique_id: str = None) -> str: - """ - It's not actually recursive - because, you know, any recursion could be turned into iteration, - and good boys do that. 😊 - """ - assert isinstance(_RECURSION_LIMIT, int) and _RECURSION_LIMIT > 0 +def _recursive_format_safe(template: str, format_dict: _t.Dict[str, _t.Any], show: bool = True, unique_id: str = None, use_safe_mode: bool = True) -> str: + """ + Safe recursive format that won't break on unknown curly brackets. + """ + assert isinstance(_RECURSION_LIMIT, int) and _RECURSION_LIMIT > 0 - prev: str = '' - new: str = template - for i in range(_RECURSION_LIMIT): - if prev == new: - break - prev = new - new = new.format_map(format_dict) - if prev == new: - return new + prev: str = '' + new: str = template + + for i in range(_RECURSION_LIMIT): + if prev == new: + break + prev = new + + if use_safe_mode: + new = _safe_format(new, format_dict) + else: + # Use the escape approach for compatibility with original behavior + new = _escape_unknown_brackets(new, format_dict) + + if prev == new: + return new - msg = ( - f"Recursion limit ({_RECURSION_LIMIT}) reached on attempt to format a string: {template!r}\n" - f"Last two formatting attempts:\n{prev!r}\n{new!r}" - ) - if show and unique_id: - _show_text_on_node(msg, unique_id) - raise RecursionError(msg) - # noinspection PyUnreachableCode - return '' # just to be safe - -# -------------------------------------- - -_dict = dict + msg = ( + f"Recursion limit ({_RECURSION_LIMIT}) reached on attempt to format a string: {template!r}\n" + f"Last two formatting attempts:\n{prev!r}\n{new!r}" + ) + if show and unique_id: + _show_text_on_node(msg, unique_id) + raise RecursionError(msg) +# Updated input types with new sanitization option _input_types = _deepfreeze({ - 'required': { - 'template': (_IO.STRING, {'multiline': True, 'tooltip': ( - "Type the text template. " - "To reference named substrings from format-dictionary, use this syntax: {substring_name}. For example:\n\n" - "score_9, score_8_up, score_7_up, {char1_short}, standing next to {char2_short},\n" - "{char1_long}\n{char2_long}" - )}), - 'recursive_format': (_IO.BOOLEAN, {'default': False, 'label_on': '❗ yes', 'label_off': 'no', 'tooltip': ( - "Do recursive format - i.e., allow the chunks from the dictionary to reference other chunks - which in turn " - "lets you do really crazy stuff like building entire HIERARCHIES of sub-prompts, all linked to the same " - "wording typed in one place.\n\n" - "The exclamation mark reminds you that with great power comes great responsibility!\n" - "You can end up with chunks cross-referencing each other in a loop. In such case, the node will just error out " - "and won't let you crash the entire ComfyUI - but still, it's your responsibility to prevent " - "such looping dictionaries." - )}), - 'show_status': (_IO.BOOLEAN, {'default': True, 'label_on': 'formatted string', 'label_off': 'no', 'tooltip': ( - "Show the final string constructed from the text-template and format-dictionary?" - )}), - }, - 'optional': { - 'dict': _DataTypes.input_dict(tooltip=( - "The dictionary to take named sub-strings from. It could be left unconnected, if the pattern doesn't reference " - "any sub-strings - then, this node acts exactly the same as a regular string-primitive node." - )), - }, - 'hidden': { - 'unique_id': 'UNIQUE_ID', # used for text display at the bottom of the node - }, + 'required': { + 'template': (_IO.STRING, {'multiline': True, 'tooltip': ( + "Type the text template. " + "To reference named substrings from format-dictionary, use this syntax: {substring_name}. For example:\n\n" + "score_9, score_8_up, score_7_up, {char1_short}, standing next to {char2_short},\n" + "{char1_long}\n{char2_long}\n\n" + "Curly brackets that don't match dictionary keys will be handled according to the sanitization mode." + )}), + 'recursive_format': (_IO.BOOLEAN, {'default': False, 'label_on': '❗ yes', 'label_off': 'no', 'tooltip': ( + "Do recursive format - i.e., allow the chunks from the dictionary to reference other chunks." + )}), + 'safe_mode': (_IO.BOOLEAN, {'default': True, 'label_on': 'safe', 'label_off': 'escape', 'tooltip': ( + "Safe mode: Only format known variables, leave unknown {brackets} as-is.\n" + "Escape mode: Escape unknown {brackets} to {{brackets}} then format normally.\n" + "Safe mode is recommended for templates with JSON, CSS, or other literal curly brackets." + )}), + 'show_status': (_IO.BOOLEAN, {'default': True, 'label_on': 'formatted string', 'label_off': 'no', 'tooltip': ( + "Show the final string constructed from the text-template and format-dictionary?" + )}), + }, + 'optional': { + 'dict': _DataTypes.input_dict(tooltip=( + "The dictionary to take named sub-strings from. It could be left unconnected, if the pattern doesn't reference " + "any sub-strings - then, this node acts exactly the same as a regular string-primitive node." + )), + }, + 'hidden': { + 'unique_id': 'UNIQUE_ID', + }, }) - class StringConstructorFormatter: - """ - Construct the formatted string from template and format-dictionary. - """ - NODE_NAME = 'StringConstructorFormatter' - CATEGORY = _meta.category - DESCRIPTION = _format_docstring(_cleandoc(__doc__)) + """ + Construct the formatted string from template and format-dictionary. + Now with safe handling of curly brackets that aren't format variables. + """ + NODE_NAME = 'StringConstructorFormatter' + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) - OUTPUT_NODE = True + OUTPUT_NODE = True - FUNCTION = 'main' - RETURN_TYPES = (_IO.STRING, ) - RETURN_NAMES = ('string', ) - # OUTPUT_TOOLTIPS = tuple() + FUNCTION = 'main' + RETURN_TYPES = (_IO.STRING, ) + RETURN_NAMES = ('string', ) - @classmethod - def INPUT_TYPES(cls): - return _input_types + @classmethod + def INPUT_TYPES(cls): + return _input_types - @staticmethod - def main( - template: str, recursive_format: bool = False, show_status: bool = False, dict: _t.Dict[str, _t.Any] = None, - unique_id: str = None - ) -> _t.Tuple[str]: - if dict is None: - dict = _dict() - _verify_input_dict(dict) - if not isinstance(template, str): - raise TypeError(f"Not a string: {template!r}") + @staticmethod + def main( + template: str, + recursive_format: bool = False, + safe_mode: bool = True, + show_status: bool = False, + dict: _t.Dict[str, _t.Any] = None, + unique_id: str = None + ) -> _t.Tuple[str]: + + if dict is None: + dict = {} + _verify_input_dict(dict) + + if not isinstance(template, str): + raise TypeError(f"Not a string: {template!r}") - out_text = ( - _recursive_format(template, dict, show_status, unique_id) - if recursive_format - else template.format_map(dict) - ) - if show_status and unique_id: - _show_text_on_node(out_text, unique_id) - return (out_text, ) + if recursive_format: + out_text = _recursive_format_safe(template, dict, show_status, unique_id, safe_mode) + else: + if safe_mode: + out_text = _safe_format(template, dict) + else: + out_text = _escape_unknown_brackets(template, dict) + + if show_status and unique_id: + _show_text_on_node(out_text, unique_id) + + return (out_text, )