diff --git a/__init__.py b/__init__.py index ed2ba9b..afa9787 100644 --- a/__init__.py +++ b/__init__.py @@ -11,9 +11,21 @@ from .nodes_extract import * class DictToolsExtension(_ComfyExtension): async def get_node_list(self) -> list[type[_io.ComfyNode]]: return [ - DictFromText, DictFromTextOld1, - DictAddAny, DictAddString, DictExtractString, - DictAddAnyOld1, DictAddStringOld1, DictExtractStringOld1 + DictAddAny, + DictAddBool, DictAddFloat, DictAddInt, DictAddString, + DictAddCond, DictAddImage, DictAddLatent, DictAddMask, + DictAddGuider, DictAddNoise, DictAddSampler, DictAddSigmas, + DictAddClip, DictAddClipVision, DictAddControlNet, DictAddGligen, DictAddLora, DictAddModel, DictAddUpscale, DictAddVae, + DictAddAnyOld1, DictAddStringOld1, + + DictExtractAny, + DictExtractBool, DictExtractFloat, DictExtractInt, DictExtractString, + DictExtractCond, DictExtractImage, DictExtractLatent, DictExtractMask, + DictExtractGuider, DictExtractNoise, DictExtractSampler, DictExtractSigmas, + DictExtractClip, DictExtractClipVision, DictExtractControlNet, DictExtractGligen, DictExtractLora, DictExtractModel, DictExtractUpscale, DictExtractVae, + DictExtractStringOld1, + + TextToDict, TextToDictOld1, TextToDictOld2, ] async def comfy_entrypoint() -> _ComfyExtension: diff --git a/__meta.py b/__meta.py index 0944892..6e9a716 100644 --- a/__meta.py +++ b/__meta.py @@ -5,7 +5,12 @@ Metadata-related module. category: str = "πŸ—‚οΈ Dict Tools" category_add: str = f"{category}/Add" +category_add_adv: str = f"{category_add}/Advanced" +category_add_mdl: str = f"{category_add}/Models" + category_extract: str = f"{category}/Extract" +category_extract_adv: str = f"{category_extract}/Advanced" +category_extract_mdl: str = f"{category_extract}/Models" category_old: str = f"{category}/_old" diff --git a/__typing.py b/__typing.py index a310bd2..5a275b5 100644 --- a/__typing.py +++ b/__typing.py @@ -3,8 +3,97 @@ """ import typing as _t -from typing import Any as _A, Optional as _O, Union as _U, TypeVar as _TypeVar +from typing import ( + Any as _A, Callable as _C, Optional as _O, Union as _U, + TypeVar as _TypeVar +) + +from sys import float_info as _float_info + +from frozendict import frozendict as _frozendict + +from comfy_api.latest import io as _io +# Not used in this module per se, but used for input-type checking: +from comfy.clip_vision import ClipVisionModel as _ClipVisionModel +from comfy.controlnet import ControlNet as _ControlNet +from comfy.model_patcher import ModelPatcher as _ModelPatcher +from comfy.samplers import CFGGuider as _CFGGuider, Sampler as _Sampler +from comfy.sd import CLIP as _CLIP, VAE as _VAE +from spandrel import ImageModelDescriptor as _ImageModelDescriptor T = _TypeVar('T') T2 = _TypeVar('T2') DictMap = _TypeVar('DictMap', bound=_U[_t.Dict, _t.Mapping]) +T_ComfyNode = _TypeVar('T_ComfyNode', bound=_io.ComfyNode) + +INT_MAX: int = 9223372036854775807 +FLOAT_MAX: float = _float_info.max + + +@_t.runtime_checkable +class NoiseProtocol(_t.Protocol): + """The expected type of ``Noise`` input.""" + # Seed should be only an int - but just to be safe: + seed: _U[int, float, str] + + def generate_noise(self, *args, **kwargs) -> _A: pass + + +# noinspection PyShadowingBuiltins +def _validate_type_factory( + type: _t.Type[T], type_name: str = None, what: str = 'Value', a='a' +) -> _C[[_A], T]: + """Build a function to type-check an input against a specific type.""" + if not what: + what = 'Value' + what = str(what) + what = what[0].upper() + what[1:] # Make the first letter uppercase + + if not a: + a = 'a' + a = str(a) + + if not type_name: + type_name = type.__name__ + type_name = str(type_name) + + error_template = f"{what} isn't {a} <{type_name}> instance: {{}}" + + def validate_type(value: _A) -> T: + """Type-check an input.""" + if not isinstance(value, type): + raise TypeError(error_template.format(repr(value))) + return value + + return validate_type + + +def _validate_required_dict(dict: _A) -> DictMap: + if dict is None: + raise ValueError("No Dict provided.") + if not isinstance(dict, (_t.Dict, _frozendict, _t.Mapping)): + raise TypeError(f"Not a Dict: {dict!r}") + if not dict: + raise ValueError("Dict is empty.") + return dict + + +def _validate_str_key(key: _A, strip: bool = False) -> str: + if key is None: + raise ValueError("No key provided.") + key = str(key) + + if not strip: + return key + + key_raw = key + key = '' + for line in key_raw.split(): + line = line.strip() # just in case \r is left + if line: + key = line + break + return key + + +_validate_str_value = _validate_type_factory(str, type_name='string') diff --git a/_io_custom.py b/_io_custom.py index 777c184..639e01d 100644 --- a/_io_custom.py +++ b/_io_custom.py @@ -7,7 +7,10 @@ from dataclasses import dataclass as _dataclass, fields as __dataclass_fields from comfy_api.latest import io as _io from .__meta import category_old as _category_old -from .__typing import _t, T as _T, _A, _O, _U, _TypeVar +from .__typing import ( + _t, _A, _C, _O, _U, _TypeVar, + T as _T, T_ComfyNode as _T_ComfyNode +) _DICT = _io.Custom("DICT") @@ -16,14 +19,31 @@ _DICT_INPUT_OPTIONAL = _DICT.Input( optional=True, tooltip="An (optional) Dictionary to work with." ) +_DICT_INPUT_REQUIRED = _DICT.Input( + 'dict', + tooltip="A Dictionary to work with." +) _DICT_OUTPUT = _DICT.Output( 'DICT', display_name='dict' ) +_KEY_CLEANUP_INPUT = _io.Boolean.Input( + 'cleanup_key', display_name='clean key', + tooltip=( + "Automatically remove leading/trailing spaces and extra newlines from the key." + ), + default=True, + label_on='strip spaces', + label_off='no', +) _KEY_INPUT_ADD = _io.String.Input( 'key', tooltip="Key (name) of the item inserted into the dict.", ) +_KEY_INPUT_EXTRACT = _io.String.Input( + 'key', + tooltip="Key (name) of the item extracted from the dict.", +) _T_Input = _t.TypeVar('T_Input', bound=_io.Input) @@ -142,3 +162,40 @@ class _BaseNode(_io.ComfyNode): _T_BaseNode = _TypeVar('T_BaseNode', bound=_BaseNode) + + +def _attach_execute_method(cls: _t.Type[_T_ComfyNode], execute: _C, docstring: str = None) -> _t.Type[_T_ComfyNode]: + """Attach a function as ``execute`` method override.""" + if not docstring: + try: + docstring = execute.__doc__ + except AttributeError: + docstring = None + if not docstring: + try: + docstring = cls.execute.__doc__ + except AttributeError: + docstring = None + if docstring: + execute.__doc__ = docstring + + # Mimic the func's metadata to make it look as if it was + # actually defined as an in-class method: + execute.__module__ = cls.__module__ + execute.__name__ = 'execute' + execute.__qualname__ = f"{cls.__qualname__}.execute" + + # Make it a class method (has to be done after this ^) + # and attach to the class: + cls.execute = classmethod(execute) + + # Since any decorator works after ABCMeta has done its job, + # and `execute` might be previously missing (kept as abstract method), + # we need to manually exclude it: + if getattr(cls, "__abstractmethods__", None): + cls.__abstractmethods__ = frozenset( + name for name in cls.__abstractmethods__ + if name != 'execute' + ) + + return cls diff --git a/node_dict_from_text.py b/node_dict_from_text.py index d683bf7..e156a7d 100644 --- a/node_dict_from_text.py +++ b/node_dict_from_text.py @@ -82,27 +82,34 @@ def _parsed_kv_pairs_gen( # ---------------------------------------------------------- -class DictFromText(_BaseNode): - """ - Build a dict of named sub-strings to be used later in string formatting (text construction). - """ +class TextToDict(_BaseNode): + """Parse raw text into key-value pairs.""" _schema = _io.Schema( - node_id=f'DictFromText{_pack_id}', - display_name='Dict from Text', + node_id=f'TextToDict{_pack_id}', + display_name='πŸ“ƒText β†’ DictπŸ—‚οΈ', category=_category, description=_format_docstring(_cleandoc(__doc__)), inputs=[ _io.Boolean.Input( 'pre_cleanup', display_name='cleanup', - tooltip="When enabled, each line is stripped from any leading/trailing spaces before parsing the text.", + tooltip=( + "When enabled, each line is individually pre-stripped " + "of any leading/trailing spaces before parsing the text." + ), default=True, - label_on='leading/trailing spaces', + label_on='strip spaces', label_off='no', ), _io.String.Input( 'dict_text', display_name='dict-items text', tooltip=( - "Sub-string names followed by their text. Different sub-string chunks are separated by empty lines. Example:\n\n" + "A wall of key-value pairs:\n" + "β€’ keys (item names) on their own line,\n" + "β€’ followed by their text.\n\n" + "Different items are separated by empty lines." + ), + placeholder=( + "Example:\n\n" "char1_short\n1boy, blond, short hair\n\n" "char1_long\n1boy, smiling, blue eyes, blond, short hair,\nwearing a leather jacket, sitting on a bike" ), @@ -125,10 +132,11 @@ class DictFromText(_BaseNode): @classmethod def execute(cls, - pre_cleanup: bool, - dict_text: str = None, show_status: bool = False, + pre_cleanup: bool, dict_text: _O[str], + show_status: bool = False, dict: _O[_DictMap] = None ) -> _io.NodeOutput: + """Parse raw text into key-value pairs.""" parsed_items = _parsed_kv_pairs_gen(dict_text, pre_strip_lines=pre_cleanup) new_dict = {k: v for k, v in parsed_items} @@ -151,10 +159,10 @@ class DictFromText(_BaseNode): # ========================================================== # Deprecated node with old ID (for backwards compatibility) -class DictFromTextOld1(_BaseNode): +class TextToDictOld1(_BaseNode): # noinspection PyProtectedMember _schema = _schema_old_node( - DictFromText._schema, + TextToDict._schema, 'StringConstructorDictFromText', inputs_converter=_InputsConverter( preserved=('pre_cleanup', 'dict_text', 'show_status', 'dict'), @@ -170,7 +178,13 @@ class DictFromTextOld1(_BaseNode): cleanup: bool, strings: str, show_status: bool = False, dict: _O[_DictMap] = None ) -> _io.NodeOutput: - return DictFromText.execute( + return TextToDict.execute( pre_cleanup=cleanup, dict_text=strings, show_status=show_status, dict=dict ) + +class TextToDictOld2(TextToDict): + _schema = _schema_old_node( + TextToDict._schema, + f'DictFromText{_pack_id}', + ) diff --git a/nodes_add.py b/nodes_add.py index 43719a4..0566564 100644 --- a/nodes_add.py +++ b/nodes_add.py @@ -9,79 +9,332 @@ from comfy_api.latest import io as _io, ui as _ui from .__meta import ( category_add as _category_add, + category_add_adv as _category_add_adv, + category_add_mdl as _category_add_mdl, pack_id_suffix as _pack_id ) -from .__typing import _A, _U, _O, _t, T as _T, DictMap as _DictMap +from .__typing import ( + _t, _A, _O, _U, _validate_type_factory, _validate_str_key, + T as _T, INT_MAX as _INT_MAX, FLOAT_MAX as _FLOAT_MAX, + DictMap as _DictMap, NoiseProtocol as _NoiseProtocol, + _ClipVisionModel, _ControlNet, _ModelPatcher, _CFGGuider, _Sampler, + _CLIP, _VAE, _ImageModelDescriptor +) from ._dict_funcs import _new_updated_dict from ._io_custom import ( - _BaseNode, - _DICT_INPUT_OPTIONAL, _DICT_OUTPUT, _KEY_INPUT_ADD, - _InputsConverter, _schema_old_node + _BaseNode, _T_BaseNode, + _DICT_INPUT_OPTIONAL, _DICT_OUTPUT, _KEY_CLEANUP_INPUT, _KEY_INPUT_ADD, + _InputsConverter, _schema_old_node, _attach_execute_method ) from .docstring_formatter import format_docstring as _format_docstring # ---------------------------------------------------------- +def _adding_node( + type_name: str, + value_type: _t.Type[_T], + value_input_cls: _t.Type[_io.Input], + prefix: str = '', + a='a', # a/an for docstring + convert_type: bool = False, + category=_category_add, + **input_kwargs +): + """Get class decorator building a ``DictAdd*`` node for a specific type.""" + type_name = str(type_name) + assert issubclass(value_input_cls, _io.Input) + if not prefix: + prefix = '' + if not a: + a = 'a' + + def class_decorator(cls: _t.Type[_T_BaseNode]): + docstring = f"Add/update {a} {type_name} item to a Dict." + if not cls.__doc__: + cls.__doc__ = docstring + + class_name = str(cls.__name__).split('.')[-1] + tooltip = f"The actual {type_name}-type item to add into the Dict." + cls._schema = _io.Schema( + node_id=f'{class_name}{_pack_id}', + display_name=f'{prefix}{type_name} β†’ DictπŸ—‚οΈ', + category=category, + description=_format_docstring(_cleandoc(cls.__doc__)), + inputs=[ + _KEY_INPUT_ADD, + _KEY_CLEANUP_INPUT, + _DICT_INPUT_OPTIONAL, + value_input_cls('value', tooltip=tooltip, **input_kwargs), + ], + outputs=[_DICT_OUTPUT], + ) + + _type_convert = value_type if convert_type else _validate_type_factory(type=value_type, type_name=type_name, a=a) + + def execute( + cls: _t.Type[_T_BaseNode], + key: str, value: _O[_T] = None, cleanup_key: bool = True, + dict: _O[_DictMap] = None, + ) -> _io.NodeOutput: + if value is None and not convert_type: + return _io.NodeOutput( + _frozendict() if dict is None else dict + ) + value = _type_convert(value) + key = _validate_str_key(key, strip=cleanup_key) + result = _new_updated_dict(dict, {key: value}) + return _io.NodeOutput(result) + + cls = _attach_execute_method(cls, execute, docstring) + + return cls + + return class_decorator + + +# ---------------------------------------------------------- +# Basic types + + +# noinspection PyAbstractClass +@_adding_node( + 'bool', _io.Boolean.Type, _io.Boolean.Input, convert_type=True, + default=False +) +class DictAddBool(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'float', _io.Float.Type, _io.Float.Input, convert_type=True, + default=0.0, min=-_FLOAT_MAX, max=_FLOAT_MAX, step=0.01, +) +class DictAddFloat(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'int', _io.Int.Type, _io.Int.Input, convert_type=True, a='an', + default=0, min=-_INT_MAX, max=_INT_MAX, step=1, + control_after_generate=False, +) +class DictAddInt(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Cond', list, _io.Conditioning.Input, prefix='🟠', + optional=True, +) +class DictAddCond(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Image', _io.Image.Type, _io.Image.Input, a='an', prefix='πŸ”΅', + optional=True, +) +class DictAddImage(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Latent', _U[_t.Dict, _t.Mapping], _io.Latent.Input, prefix='🟣', + optional=True, +) +class DictAddLatent(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Mask', _io.Mask.Type, _io.Mask.Input, prefix='🟒', + optional=True, +) +class DictAddMask(_BaseNode): pass + + +# ---------------------------------------------------------- +# Advanced types + + +# noinspection PyAbstractClass +@_adding_node( + 'Guider', _CFGGuider, _io.Guider.Input, + category=_category_add_adv, + optional=True, +) +class DictAddGuider(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Noise', _NoiseProtocol, _io.Noise.Input, + category=_category_add_adv, + optional=True, +) +class DictAddNoise(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Sampler', _Sampler, _io.Sampler.Input, + category=_category_add_adv, + optional=True, +) +class DictAddSampler(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Sigmas', _io.Sigmas.Type, _io.Sigmas.Input, + category=_category_add_adv, + optional=True, +) +class DictAddSigmas(_BaseNode): pass + + +# ---------------------------------------------------------- +# Model types + + +# noinspection PyAbstractClass +@_adding_node( + 'CLIP', _CLIP, _io.Clip.Input, prefix='🟑', + category=_category_add_mdl, + optional=True, +) +class DictAddClip(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'CLIP-vision', _ClipVisionModel, _io.ClipVision.Input, prefix='πŸ”΅', + category=_category_add_mdl, + optional=True, +) +class DictAddClipVision(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'ControlNet', _ControlNet, _io.ControlNet.Input, prefix='🟒', + category=_category_add_mdl, + optional=True, +) +class DictAddControlNet(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'GLIGEN', _ModelPatcher, _io.Gligen.Input, prefix='🟀', + category=_category_add_mdl, + optional=True, +) +class DictAddGligen(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'LoRA', _io.LoraModel.Type, _io.LoraModel.Input, prefix='βšͺ', + category=_category_add_mdl, + optional=True, +) +class DictAddLora(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Model', _ModelPatcher, _io.Model.Input, prefix='🟣', + category=_category_add_mdl, + optional=True, +) +class DictAddModel(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'Upscale', _ImageModelDescriptor, _io.UpscaleModel.Input, a='an', prefix='🟒', + category=_category_add_mdl, + optional=True, +) +class DictAddUpscale(_BaseNode): pass + + +# noinspection PyAbstractClass +@_adding_node( + 'VAE', _VAE, _io.Vae.Input, prefix='πŸ”΄', + category=_category_add_mdl, + optional=True, +) +class DictAddVae(_BaseNode): pass + + +# ========================================================== +# Nodes with custom implementation + + class DictAddAny(_BaseNode): - """Add/update any-type item to a Format-Dict.""" + """Add/update any-type item to a Dict.""" _schema = _io.Schema( node_id=f'DictAddAny{_pack_id}', - display_name='ANY to Dict', + display_name='β—―ANY β†’ DictπŸ—‚οΈ', category=_category_add, description=_format_docstring(_cleandoc(__doc__)), inputs=[ _KEY_INPUT_ADD, + _KEY_CLEANUP_INPUT, + _DICT_INPUT_OPTIONAL, _io.AnyType.Input( 'value', optional=True, - tooltip="The actual any-type item to add into the dict." + tooltip="The actual any-type item to add into the Dict." ), - _DICT_INPUT_OPTIONAL, ], outputs=[_DICT_OUTPUT], ) + # noinspection PyShadowingBuiltins @classmethod def execute(cls, - key: _O[str], value: _A = None, dict: _O[_DictMap] = None, + key: str, value: _A = None, cleanup_key: bool = True, + dict: _O[_DictMap] = None, ) -> _io.NodeOutput: - """Update/append an item of any type to the dict.""" - if key is None: - if dict is None: - dict = _frozendict() - # No need to create another dict instance if we add nothing: - result = dict - else: - result = _new_updated_dict(dict, {key: value}) + """Add/update any-type item to a Dict.""" + key = _validate_str_key(key, strip=cleanup_key) + # No value-type check + result = _new_updated_dict(dict, {key: value}) return _io.NodeOutput(result) + # ---------------------------------------------------------- + class DictAddString(_BaseNode): - """Add/update a string to a Format-Dict.""" + """Add/update a string item to a Dict.""" _schema = _io.Schema( node_id=f'DictAddString{_pack_id}', - display_name='STRING to Dict', + display_name='string β†’ DictπŸ—‚οΈ', category=_category_add, description=_format_docstring(_cleandoc(__doc__)), inputs=[ _KEY_INPUT_ADD, - _io.Boolean.Input( - 'cleanup_value', - tooltip=( - "When enabled, each line in the sub-string is stripped " - "from any spaces at its start and end." - ), - default=True, - label_on='leading/trailing spaces', - label_off='no', - ), + _KEY_CLEANUP_INPUT, _io.String.Input( 'value', - tooltip="The actual string to add into the dict.", + tooltip="The actual string to add into the Dict.", multiline=True, ), + _io.Boolean.Input( + 'cleanup_value', display_name='clean val', + tooltip=( + "When enabled, each line in the value-string is " + "individually stripped of any leading/trailing spaces." + ), + default=True, + label_on='strip spaces', + label_off='no', + ), _DICT_INPUT_OPTIONAL, ], @@ -90,10 +343,13 @@ class DictAddString(_BaseNode): @classmethod def execute(cls, - key: str, cleanup_value: bool, value: str = None, + key: str, value: str, + cleanup_key: bool = True, cleanup_value: bool = True, dict: _O[_DictMap] = None, ) -> _io.NodeOutput: - """Update/append a string to the dict.""" + """Add/update a string item to a Dict.""" + key = _validate_str_key(key, strip=cleanup_key) + value = '' if value is None else str(value) if cleanup_value: value = '\n'.join( @@ -146,5 +402,7 @@ class DictAddStringOld1(_BaseNode): dict: _O[_DictMap] = None, ) -> _io.NodeOutput: return DictAddString.execute( - key=name, cleanup_value=cleanup, value=string, dict=dict + key=name, value=string, + cleanup_key=False, cleanup_value=cleanup, + dict=dict ) diff --git a/nodes_extract.py b/nodes_extract.py index fff93ed..35e4e72 100644 --- a/nodes_extract.py +++ b/nodes_extract.py @@ -7,35 +7,299 @@ from comfy_api.latest import io as _io, ui as _ui from .__meta import ( category_extract as _category_extract, + category_extract_adv as _category_extract_adv, + category_extract_mdl as _category_extract_mdl, pack_id_suffix as _pack_id ) -from .__typing import _A, _U, _O, _t, T as _T, DictMap as _DictMap +from .__typing import ( + _t, _A, _O, _U, + _validate_type_factory, _validate_required_dict, _validate_str_key, + T as _T, DictMap as _DictMap, NoiseProtocol as _NoiseProtocol, + _ClipVisionModel, _ControlNet, _ModelPatcher, _CFGGuider, _Sampler, + _CLIP, _VAE, _ImageModelDescriptor +) from ._io_custom import ( - _BaseNode, - _DICT_INPUT_OPTIONAL, - _InputsConverter, _schema_old_node + _BaseNode, _T_BaseNode, + _DICT_INPUT_REQUIRED, _KEY_CLEANUP_INPUT, _KEY_INPUT_EXTRACT, + _InputsConverter, _schema_old_node, _attach_execute_method ) from .docstring_formatter import format_docstring as _format_docstring # ---------------------------------------------------------- -class DictExtractString(_BaseNode): - """Extract a single string from a Format-Dict.""" +# noinspection PyShadowingBuiltins +def _get_key(dict: _DictMap, key: _A) -> _A: + try: + return dict[key] + except KeyError: + raise KeyError(f"No such key in Dict: {key!r}") + +# ---------------------------------------------------------- + +def _extracting_node( + type_name: str, + value_type: _t.Type[_T], + value_output_cls: _t.Type[_io.Output], + suffix: str = '', + a='a', # a/an for docstring + convert_type: bool = False, + category=_category_extract, + **output_kwargs +): + """Get class decorator building a ``DictExtract*`` node for a specific type.""" + type_name = str(type_name) + assert issubclass(value_output_cls, _io.Output) + if not suffix: + suffix = '' + if not a: + a = 'a' + + def class_decorator(cls: _t.Type[_T_BaseNode]): + docstring = f"Extract {a} {type_name} item from a Dict." + if not cls.__doc__: + cls.__doc__ = docstring + + class_name = str(cls.__name__).split('.')[-1] + tooltip = f"The actual {type_name}-type item extracted from the Dict." + cls._schema = _io.Schema( + node_id=f'{class_name}{_pack_id}', + display_name=f'πŸ—‚οΈDict β†’ {type_name}{suffix}', + category=category, + description=_format_docstring(_cleandoc(cls.__doc__)), + inputs=[ + _DICT_INPUT_REQUIRED, + _KEY_INPUT_EXTRACT, + _KEY_CLEANUP_INPUT, + ], + outputs=[ + value_output_cls('value', tooltip=tooltip, **output_kwargs), + ], + ) + + _type_convert = value_type if convert_type else _validate_type_factory(type=value_type, type_name=type_name, a=a) + + # noinspection PyShadowingBuiltins + def execute( + cls: _t.Type[_T_BaseNode], + dict: _DictMap, key: str, cleanup_key: bool = True, + ) -> _io.NodeOutput: + dict = _validate_required_dict(dict) + key = _validate_str_key(key, strip=cleanup_key) + value = _get_key(dict, key) + value = _type_convert(value) + return _io.NodeOutput(value) + + cls = _attach_execute_method(cls, execute, docstring) + + return cls + + return class_decorator + + +# ---------------------------------------------------------- +# Basic types + + +# noinspection PyAbstractClass +@_extracting_node( + 'bool', _io.Boolean.Type, _io.Boolean.Output, convert_type=True, +) +class DictExtractBool(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'float', _io.Float.Type, _io.Float.Output, convert_type=True, +) +class DictExtractFloat(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'int', _io.Int.Type, _io.Int.Output, a='an',convert_type=True, +) +class DictExtractInt(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Cond', list, _io.Conditioning.Output, suffix='🟠', +) +class DictExtractCond(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Image', _io.Image.Type, _io.Image.Output, a='an', suffix='πŸ”΅', +) +class DictExtractImage(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Latent', _U[_t.Dict, _t.Mapping], _io.Latent.Output, suffix='🟣', +) +class DictExtractLatent(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Mask', _io.Mask.Type, _io.Mask.Output, suffix='🟒', +) +class DictExtractMask(_BaseNode): pass + + +# ---------------------------------------------------------- +# Advanced types + + +# noinspection PyAbstractClass +@_extracting_node( + 'Guider', _CFGGuider, _io.Guider.Output, + category=_category_extract_adv, +) +class DictExtractGuider(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Noise', _NoiseProtocol, _io.Noise.Output, + category=_category_extract_adv, +) +class DictExtractNoise(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Sampler', _Sampler, _io.Sampler.Output, + category=_category_extract_adv, +) +class DictExtractSampler(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Sigmas', _io.Sigmas.Type, _io.Sigmas.Output, + category=_category_extract_adv, +) +class DictExtractSigmas(_BaseNode): pass + + +# ---------------------------------------------------------- +# Model types + + +# noinspection PyAbstractClass +@_extracting_node( + 'CLIP', _CLIP, _io.Clip.Output, suffix='🟑', + category=_category_extract_mdl, +) +class DictExtractClip(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'CLIP-vision', _ClipVisionModel, _io.ClipVision.Output, suffix='πŸ”΅', + category=_category_extract_mdl, +) +class DictExtractClipVision(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'ControlNet', _ControlNet, _io.ControlNet.Output, suffix='🟒', + category=_category_extract_mdl, +) +class DictExtractControlNet(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'GLIGEN', _ModelPatcher, _io.Gligen.Output, suffix='🟀', + category=_category_extract_mdl, +) +class DictExtractGligen(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'LoRA', _io.LoraModel.Type, _io.LoraModel.Output, suffix='βšͺ', + category=_category_extract_mdl, +) +class DictExtractLora(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Model', _ModelPatcher, _io.Model.Output, suffix='🟣', + category=_category_extract_mdl, +) +class DictExtractModel(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'Upscale', _ImageModelDescriptor, _io.UpscaleModel.Output, a='an', suffix='🟒', + category=_category_extract_mdl, +) +class DictExtractUpscale(_BaseNode): pass + + +# noinspection PyAbstractClass +@_extracting_node( + 'VAE', _VAE, _io.Vae.Output, suffix='πŸ”΄', + category=_category_extract_mdl, +) +class DictExtractVae(_BaseNode): pass + + +# ========================================================== +# Nodes with custom implementation + + +class DictExtractAny(_BaseNode): + """Extract any-type item from a Dict.""" _schema = _io.Schema( - node_id=f'DictExtractString{_pack_id}', - display_name='STRING from Dict', + node_id=f'DictExtractAny{_pack_id}', + display_name='πŸ—‚οΈDict β†’ ANYβ—―', category=_category_extract, description=_format_docstring(_cleandoc(__doc__)), inputs=[ - _io.Boolean.Input( - 'cleanup_key', display_name='cleanup', - tooltip=( - "Automatically remove leading/trailing spaces and newlines from the key." - ), - default=True, - label_on='leading/trailing spaces', - label_off='no', + _DICT_INPUT_REQUIRED, + _KEY_INPUT_EXTRACT, + _KEY_CLEANUP_INPUT, + ], + outputs=[ + _io.AnyType.Output( + 'value', + tooltip="The actual any-type item extracted from the Dict." ), + ], + # hidden=[_io.Hidden.unique_id], + ) + + # noinspection PyShadowingBuiltins + @classmethod + def execute(cls, + dict: _DictMap, key: str, cleanup_key: bool = True, + ) -> _io.NodeOutput: + """Extract any-type item from a Dict.""" + dict = _validate_required_dict(dict) + key = _validate_str_key(key, strip=cleanup_key) + value = _get_key(dict, key) + # No output-type check + return _io.NodeOutput(value) + + +class DictExtractString(_BaseNode): + """Extract a string item from a Dict.""" + _schema = _io.Schema( + node_id=f'DictExtractString{_pack_id}', + display_name='πŸ—‚οΈDict β†’ string', + category=_category_extract, + description=_format_docstring(_cleandoc(__doc__)), + inputs=[ + _DICT_INPUT_REQUIRED, _io.String.Input( 'key', tooltip=( @@ -44,6 +308,7 @@ class DictExtractString(_BaseNode): "If no such key exists in the dict, an empty string returned." ), ), + _KEY_CLEANUP_INPUT, _io.Boolean.Input( 'show_status', tooltip="Show the extracted string on the node itself?", @@ -51,44 +316,34 @@ class DictExtractString(_BaseNode): label_on='value', label_off='no', ), - - _DICT_INPUT_OPTIONAL, ], - outputs=[_io.String.Output('string')], + outputs=[ + _io.String.Output( + 'value', + tooltip="The actual string item extracted from the Dict." + ) + ], # hidden=[_io.Hidden.unique_id], ) @classmethod def execute(cls, - cleanup_key: bool, key: str, show_status: bool = False, - dict: _O[_DictMap] = None + dict: _DictMap, key: str, + cleanup_key: bool = True, show_status: bool = False, ) -> _io.NodeOutput: - """Extract a single string from the Format-Dict.""" - if cleanup_key: - for line in key.split(): - line = line.strip() # just in case \r is left - if line: - key = line - break + """Extract a string item from a Dict.""" + dict = _validate_required_dict(dict) + key = _validate_str_key(key, strip=cleanup_key) + value = _get_key(dict, key) - # The passed dict might not actually be a dict but ANY mapping, - # which doesn't have a `get()` method. - # So direct key access is better: - string: str = '' - # noinspection PyBroadException - try: - string = dict[key] - except Exception: - pass - - if string is None: - string = '' - elif not isinstance(string, str): - string = repr(string) + if value is None: + value = '' + elif not isinstance(value, str): + value = repr(value) return _io.NodeOutput( - string, - ui=_ui.PreviewText(string) if show_status else None + value, + ui=_ui.PreviewText(value) if show_status else None ) @@ -113,5 +368,5 @@ class DictExtractStringOld1(_BaseNode): dict: _O[_DictMap] = None ) -> _io.NodeOutput: return DictExtractString.execute( - cleanup_key=False, key=name, show_status=show_status, dict=dict + dict=dict, key=name, cleanup_key=False, show_status=show_status ) diff --git a/pyproject.toml b/pyproject.toml index 69b0721..8793ae8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ [project] name = "dict_tools" -version = "3.1.2" +version = "3.2.0" description = "Essential nodes to use dictionaries in ComfyUI: for smart prompt-formatting, general organization (passing a single connection instead of spaghetti), or anything else." license = {file = "LICENSE.md"} readme = "README.md" authors = [ {name = "Lex Darlog"} ] -requires-python = ">=3.7" +requires-python = ">=3.8" # Protocols are used dependencies = ["frozendict"] classifiers = [ # https://docs.comfy.org/registry/specifications#classifiers-recommended