From 3a14652dc9a3365e8fd32daf12b5327a779a66c8 Mon Sep 17 00:00:00 2001 From: "Lex Darlog (DRL)" <3897975+Lex-DRL@users.noreply.github.com> Date: Sun, 8 Mar 2026 23:08:00 -0300 Subject: [PATCH] Update all the remaining dict-modifying nodes + change approach to intermediate dict -> preserve type --- src/basic_data_handling/dict_nodes.py | 192 ++++++++++++++++++-------- 1 file changed, 137 insertions(+), 55 deletions(-) diff --git a/src/basic_data_handling/dict_nodes.py b/src/basic_data_handling/dict_nodes.py index d3a72fa..eddfaa3 100644 --- a/src/basic_data_handling/dict_nodes.py +++ b/src/basic_data_handling/dict_nodes.py @@ -1,6 +1,7 @@ -from typing import Any +from typing import Any, Mapping, Type + from inspect import cleandoc -from itertools import chain +import random try: from comfy.comfy_types.node_typing import IO, ComfyNodeABC @@ -17,6 +18,24 @@ except: from ._dynamic_input import ContainsDynamicDict +def _output_dict_preserving_type(result: dict, in_type: Type[Mapping]): + if in_type is dict: + return result + + # Support other dict-like classes: + # noinspection PyBroadException + try: + return in_type(result) + except Exception: + pass + + # noinspection PyBroadException + try: + return in_type(result.items()) + except Exception: + return result + + class DictCreate(ComfyNodeABC): """ Creates a new empty dictionary. @@ -351,7 +370,30 @@ class DictExcludeKeys(ComfyNodeABC): FUNCTION = "exclude_keys" def exclude_keys(self, input_dict: dict, keys_to_exclude: list) -> tuple[dict]: - result = {k: v for k, v in input_dict.items() if k not in keys_to_exclude} + if not(input_dict and keys_to_exclude): + return (input_dict,) + + # `in` check is faster with sets + remove duplicates: + keys_to_exclude_set = set(keys_to_exclude) + + if len(input_dict) <= len(keys_to_exclude_set): + # It's faster to rebuild the dict with filter. + result = { + k: v for k, v in input_dict.items() + if k not in keys_to_exclude_set + } + else: + # It's faster to duplicate the dict, then pop all the exclusions one-by-one. + result = dict(input_dict) + for key in keys_to_exclude_set: + if key in result: + result.pop(key) + + if len(result) == len(input_dict): + # No changes made + return (input_dict,) + + result = _output_dict_preserving_type(result, type(input_dict)) return (result,) @@ -377,7 +419,30 @@ class DictFilterByKeys(ComfyNodeABC): FUNCTION = "filter_by_keys" def filter_by_keys(self, input_dict: dict, keys: list) -> tuple[dict]: - result = {k: input_dict[k] for k in keys if k in input_dict} + if not(input_dict and keys): + return (input_dict,) + + # `in` check is faster with sets + remove duplicates: + keys_set = set(keys) + + if len(keys_set) <= len(input_dict): + # It's faster to iterate over preserved keys + result = { + k: input_dict[k] for k in keys_set + if k in input_dict + } + else: + # It's faster to iterate over the input dict items + result = { + k: v for k, v in input_dict.items() + if k in keys_set + } + + if len(result) == len(input_dict): + # No changes made + return (input_dict,) + + result = _output_dict_preserving_type(result, type(input_dict)) return (result,) @@ -522,11 +587,13 @@ class DictInvert(ComfyNodeABC): def invert(self, input_dict: dict) -> tuple[dict, bool]: try: inverted = {v: k for k, v in input_dict.items()} - return inverted, True except Exception: # Return original dictionary if inversion fails (e.g., unhashable values) return input_dict, False + inverted = _output_dict_preserving_type(inverted, type(input_dict)) + return inverted, True + class DictItems(ComfyNodeABC): """ @@ -629,14 +696,14 @@ class DictMerge(ComfyNodeABC): if not extra_dicts: return (dict1,) - # dict1 might be something like frozendict or other immutable mapping class. - # Thus, we shouldn't just copy-and-update, - # but instead we should build a new instance of the same type: - dict1_type = type(dict1) - result = dict1_type(chain( - dict1.items(), - *(x.items() for x in extra_dicts), - )) + # dict1 might be something like `frozendict` or other immutable mapping class. + # Thus, we shouldn't do `dict.copy()`, + # we must explicitly construct a `dict` object: + result = dict(dict1) + for extra in extra_dicts: + result.update(extra) + + result = _output_dict_preserving_type(result, type(dict1)) return (result,) @@ -668,17 +735,21 @@ class DictPop(ComfyNodeABC): FUNCTION = "pop" def pop(self, input_dict: dict, key: str, default_value=None) -> tuple[dict, Any]: - result = input_dict.copy() + if key not in input_dict: + return input_dict, default_value + # input_dict might be something like `frozendict` or other immutable mapping class. + # Thus, we shouldn't do `dict.copy()`, + # we must explicitly construct a `dict` object: + result = dict(input_dict) try: - if key in result: - value = result.pop(key) - return result, value - else: - return result, default_value + value = result.pop(key) except Exception as e: raise ValueError(f"Error popping key from dictionary: {str(e)}") + result = _output_dict_preserving_type(result, type(input_dict)) + return result, value + class DictPopItem(ComfyNodeABC): """ @@ -686,7 +757,8 @@ class DictPopItem(ComfyNodeABC): This node takes a dictionary as input, removes an arbitrary key-value pair, and returns the modified dictionary along with the removed key and value. - If the dictionary is empty, returns an error. + If operation fails, no error is thrown, but the last argument is `False`, + and the dictionary is returned intact. """ @classmethod def INPUT_TYPES(cls): @@ -703,15 +775,18 @@ class DictPopItem(ComfyNodeABC): FUNCTION = "popitem" def popitem(self, input_dict: dict) -> tuple[dict, str, Any, bool]: - result = input_dict.copy() + if not input_dict: + return input_dict, "", None, False + + result = dict(input_dict) + # noinspection PyBroadException try: - if result: - key, value = result.popitem() - return result, key, value, True - else: - return result, "", None, False - except: - return result, "", None, False + key, value = result.popitem() + except Exception: + return input_dict, "", None, False + + result = _output_dict_preserving_type(result, type(input_dict)) + return result, key, value, True class DictPopRandom(ComfyNodeABC): @@ -741,17 +816,18 @@ class DictPopRandom(ComfyNodeABC): return float("NaN") # Not equal to anything -> trigger recalculation def pop_random(self, input_dict: dict) -> tuple[dict, str, Any, bool]: - import random - result = input_dict.copy() + if not input_dict: + return input_dict, "", None, False + + result = dict(input_dict) try: - if result: - random_key = random.choice(list(result.keys())) - random_value = result.pop(random_key) - return result, random_key, random_value, True - else: - return result, "", None, False - except: - return result, "", None, False + random_key = random.choice(list(result.keys())) + random_value = result.pop(random_key) + except Exception: + return input_dict, "", None, False + + result = _output_dict_preserving_type(result, type(input_dict)) + return result, random_key, random_value, True class DictRemove(ComfyNodeABC): @@ -778,11 +854,14 @@ class DictRemove(ComfyNodeABC): FUNCTION = "remove" def remove(self, input_dict: dict, key: str) -> tuple[dict, bool]: - result = input_dict.copy() - if key in result: - del result[key] - return result, True - return result, False + if key not in input_dict: + return input_dict, False + + result = dict(input_dict) + del result[key] + + result = _output_dict_preserving_type(result, type(input_dict)) + return result, True class DictSet(ComfyNodeABC): @@ -808,11 +887,10 @@ class DictSet(ComfyNodeABC): FUNCTION = "set" def set(self, input_dict: dict, key: str, value: Any) -> tuple[dict]: - # input_dict might be something like frozendict or other immutable mapping class. - # Thus, we shouldn't just copy-and-update, - # but instead we should build a new instance of the same type: - dict_type = type(input_dict) - result = dict_type(chain(input_dict.items(), (key, value))) + result = dict(input_dict) + result[key] = value + + result = _output_dict_preserving_type(result, type(input_dict)) return (result,) @@ -840,9 +918,11 @@ class DictSetDefault(ComfyNodeABC): DESCRIPTION = cleandoc(__doc__ or "") FUNCTION = "setdefault" - def setdefault(self, input_dict: dict, key: str, default_value=None) -> tuple[dict, Any]: - result = input_dict.copy() + def setdefault(self, input_dict: dict, key: str, default_value: Any = None) -> tuple[dict, Any]: + result = dict(input_dict) value = result.setdefault(key, default_value) + + result = _output_dict_preserving_type(result, type(input_dict)) return result, value @@ -869,11 +949,13 @@ class DictUpdate(ComfyNodeABC): FUNCTION = "update" def update(self, dict1: dict, dict2: dict) -> tuple[dict]: - # dict1 might be something like frozendict or other immutable mapping class. - # Thus, we shouldn't just copy-and-update, - # but instead we should build a new instance of the same type: - dict1_type = type(dict1) - result = dict1_type(chain(dict1.items(), dict2.items())) + if not dict2: + return (dict1,) + + result = dict(dict1) + result.update(dict2) + + result = _output_dict_preserving_type(result, type(dict1)) return (result,)