diff --git a/node_formatter.py b/node_formatter.py index ce2210d..5a7a5e3 100644 --- a/node_formatter.py +++ b/node_formatter.py @@ -1,10 +1,12 @@ # encoding: utf-8 """ +Code for ``StringConstructorFormatter`` node. """ import typing as _t from inspect import cleandoc as _cleandoc +import re as _re import sys as _sys from frozendict import deepfreeze as _deepfreeze @@ -15,14 +17,60 @@ 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) # You can externally monkey-patch it... but if it blows up, your fault 🤷🏻‍♂️ -def _recursive_format(template: str, format_dict: _t.Dict[str, _t.Any], show: bool = True, unique_id: str = None) -> str: +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_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. + It's not actually recursive - because, you know, any recursion could be turned into iteration, and good boys do that. 😊 """ @@ -34,7 +82,13 @@ def _recursive_format(template: str, format_dict: _t.Dict[str, _t.Any], show: bo if prev == new: break prev = new - new = new.format_map(format_dict) + + 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 @@ -50,24 +104,22 @@ def _recursive_format(template: str, format_dict: _t.Dict[str, _t.Any], show: bo # -------------------------------------- -_dict = dict - _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}" + "{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 - 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." + "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?" @@ -106,20 +158,27 @@ class StringConstructorFormatter: @staticmethod def main( - template: str, recursive_format: bool = False, show_status: bool = False, dict: _t.Dict[str, _t.Any] = None, + 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 = _dict() + 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 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, )