239 lines
8.6 KiB
Python
239 lines
8.6 KiB
Python
# encoding: utf-8
|
|
"""
|
|
Code for ``StringConstructorFormatter`` node.
|
|
"""
|
|
|
|
import typing as _t
|
|
|
|
from dataclasses import dataclass as _dataclass
|
|
from inspect import cleandoc as _cleandoc
|
|
import re as _re
|
|
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
|
|
|
|
|
|
_RECURSION_LIMIT = max(int(_sys.getrecursionlimit()), 1) # You can externally monkey-patch it... but if it blows up, your fault 🤷🏻♂️
|
|
|
|
def _n_brackets_after_escape_for_existing_key(n_brackets: int):
|
|
"""
|
|
Calculate the updated number of braces before/after the pattern.
|
|
Keeps the innermost set of brackets as-is for formatting,
|
|
but doubles all the leading/trailing brackets to escape them.
|
|
"""
|
|
n_extra = max(n_brackets - 1, 0)
|
|
return n_extra * 2 + 1
|
|
|
|
|
|
def _rebuild_parsed_keyword(n_opening_brackets: int, inside_brackets: str, n_closing_brackets: int) -> str:
|
|
# According to "Beautiful/ideomatic python" lecture, string formatting is faster than other forms of concatenation
|
|
return '{}{}{}'.format(
|
|
'{' * n_opening_brackets,
|
|
inside_brackets,
|
|
'}' * n_closing_brackets
|
|
)
|
|
|
|
|
|
_re_formatting_keyword_sub = _re.compile( # Pre-compiled regex-sub func to find and replace {keyword} patterns
|
|
'(\{+)([^{}]*)(\}+)'
|
|
# TODO: Takes leading/trailing braces into account, but not nested ones. Maybe implement some day,
|
|
# but it would require true parsing - with stack of braces, nested parts iterator, etc.
|
|
).sub
|
|
|
|
|
|
@_dataclass
|
|
class _Formatter:
|
|
"""
|
|
A callable (function-like) class, which does the actual formatting, while respecting all the options.
|
|
|
|
It's made as a class to split the formatting into two stages:
|
|
- First, the instance is properly initialized with the shared arguments (all the method overrides are made depending on options);
|
|
- Then, the actual instance is treated as a function - it needs to be called with the formatted string as the only argument.
|
|
|
|
It's done this way to avoid extra conditions in the loop + to organize all the setup work in one place.
|
|
"""
|
|
format_dict: _t.Optional[_t.Dict[str, _t.Any], None] = None
|
|
|
|
recursive: bool = False
|
|
safe: bool = True
|
|
|
|
show_status: bool = True
|
|
unique_node_id: str = None
|
|
|
|
def __post_init__(self): # called by dataclass init
|
|
if self.format_dict is None:
|
|
self.format_dict = dict()
|
|
|
|
format_dict = self.format_dict
|
|
_verify_input_dict(format_dict)
|
|
if not format_dict:
|
|
self.__escape_match = self.__escape_match_no_keys
|
|
|
|
# No need to override _escape_for_safe_format().
|
|
# Instead, recursive method defines local function, and simple method just checks self.safe mode in-place.
|
|
|
|
self._format = self.__format_recursive if self.recursive else self.__format_simple
|
|
|
|
# @staticmethod
|
|
# def __dummy_return_intact(template: str) -> str:
|
|
# return template
|
|
|
|
def __escape_match(self, match: _re.Match[str]) -> str:
|
|
opening_brackets, inside_brackets, closing_brackets = match.groups()
|
|
if inside_brackets.strip() in self.format_dict:
|
|
n_opening = _n_brackets_after_escape_for_existing_key(opening_brackets)
|
|
n_closing = _n_brackets_after_escape_for_existing_key(closing_brackets)
|
|
return _rebuild_parsed_keyword(n_opening, inside_brackets, n_closing)
|
|
else:
|
|
return _rebuild_parsed_keyword(len(opening_brackets) * 2, inside_brackets, len(closing_brackets) * 2)
|
|
|
|
@staticmethod
|
|
def __escape_match_no_keys(match: _re.Match[str]) -> str:
|
|
opening_brackets, inside_brackets, closing_brackets = match.groups()
|
|
return _rebuild_parsed_keyword(len(opening_brackets) * 2, inside_brackets, len(closing_brackets) * 2)
|
|
|
|
def _escape_for_safe_format(self, template: str) -> str:
|
|
"""
|
|
Method to be called in safe mode - on template, before the actual format.
|
|
|
|
Safe mode only replaces variables that exist in the dictionary. Leaves other curly brackets untouched.
|
|
"""
|
|
return _re_formatting_keyword_sub(self.__escape_match, template)
|
|
|
|
def __format_simple(self, template: str) -> str:
|
|
if self.safe:
|
|
template = self._escape_for_safe_format(template)
|
|
return template.format_map(self.format_dict)
|
|
|
|
def __format_recursive(self, template: str) -> 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
|
|
|
|
# Cache before loop - to avoid even 'dot' operator
|
|
escape_func = self._escape_for_safe_format
|
|
format_dict: _t.Dict[str, _t.Any] = self.format_dict
|
|
|
|
# The following two funcs defined here instead of method overrides - to simplify their local scope:
|
|
|
|
def format_single_with_escape(_template: str):
|
|
_template = escape_func(_template)
|
|
return _template.format_map(format_dict)
|
|
|
|
def format_single_no_escape(_template: str):
|
|
return _template.format_map(format_dict)
|
|
|
|
format_single_func = format_single_with_escape if self.safe else format_single_no_escape
|
|
|
|
prev: str = ''
|
|
new: str = template
|
|
for i in range(_RECURSION_LIMIT):
|
|
if prev == new:
|
|
return new
|
|
prev = new
|
|
new = format_single_func(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 self.show_status and self.unique_node_id:
|
|
_show_text_on_node(msg, self.unique_node_id)
|
|
raise RecursionError(msg)
|
|
# noinspection PyUnreachableCode
|
|
return '' # just to be extra-safe, if RecursionError is treated as warning
|
|
|
|
@staticmethod
|
|
def _format(template: str) -> str:
|
|
raise NotImplementedError("This method should always be overridden after initialization.")
|
|
|
|
def __call__(self, template: str) -> str:
|
|
if not isinstance(template, str):
|
|
raise TypeError(f"Not a string: {template!r}")
|
|
|
|
out_text = self._format(template) if template else ''
|
|
if self.show_status and self.unique_node_id:
|
|
_show_text_on_node(out_text, self.unique_node_id)
|
|
return out_text
|
|
|
|
# --------------------------------------
|
|
|
|
_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}\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_format': (_IO.BOOLEAN, {'default': True, 'label_on': 'yes', 'label_off': 'no', 'tooltip': (
|
|
"Safe mode: If a specific {keyword} doesn't exist in the dict, leave it as-is.\n"
|
|
"On: missing {keywords} preserved intact.\n"
|
|
"Off: missing {keywords} raise an error.\n\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=( # It's not actually optional, but is here since there's no dict-widget
|
|
"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
|
|
},
|
|
})
|
|
|
|
|
|
class StringConstructorFormatter:
|
|
"""
|
|
Construct the formatted string from template and format-dictionary.
|
|
"""
|
|
NODE_NAME = 'StringConstructorFormatter'
|
|
CATEGORY = _meta.category
|
|
DESCRIPTION = _format_docstring(_cleandoc(__doc__))
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = 'main'
|
|
RETURN_TYPES = (_IO.STRING, )
|
|
RETURN_NAMES = ('string', )
|
|
# OUTPUT_TOOLTIPS = tuple()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return _input_types
|
|
|
|
@staticmethod
|
|
def main(
|
|
template: str,
|
|
recursive_format: bool = False,
|
|
safe_format: bool = True,
|
|
show_status: bool = False,
|
|
dict: _t.Dict[str, _t.Any] = None, #actually, required - but it's here to keep the declared params order
|
|
unique_id: str = None
|
|
) -> _t.Tuple[str]:
|
|
formatter = _Formatter(
|
|
format_dict=dict,
|
|
recursive=recursive_format, safe=safe_format,
|
|
show_status=show_status, unique_node_id=unique_id,
|
|
)
|
|
out_text = formatter(template)
|
|
return (out_text, )
|