Merge pull request #12 from StableLlama/collections_enhancement

Collections enhancement

Author: StableLlama

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "basic_data_handling"
-version = "0.3.3"
+version = "0.3.4"
 description = """NOTE: Still in development! Expect breaking changes!
 Basic Python functions for manipulating data that every programmer is used to.
 Currently supported ComfyUI data types: BOOLEAN, FLOAT, INT, STRING and data lists.
@@ -52,9 +52,9 @@ pythonpath = [
 testpaths = [
     "tests",
 ]
-#python_files = ["test_*.py"]
+python_files = ["test_*.py"]
 #python_files = ["conftest.py", "test_boolean_nodes.py"]
-python_files = ["test_boolean_nodes.py"]
+#python_files = ["test_boolean_nodes.py"]
 
 [tool.mypy]
 files = "."

src/basic_data_handling/_dynamic_input.py

@@ -24,7 +24,6 @@ def __contains__(self, key):
 
     def __getitem__(self, key):
         # Dynamically return the value for keys matching a `prefix<number>` pattern
-        print(f'_ dynamic prefixes: {self._dynamic_prefixes}; get key: {key}')
         for prefix, value in self._dynamic_prefixes.items():
             if key.startswith(prefix) and key[len(prefix):].isdigit():
                 return value

src/basic_data_handling/data_list_nodes.py

@@ -348,6 +348,33 @@ def select(self, **kwargs: list[Any]) -> tuple[list[Any]]:
         return result_true, result_false
 
 
+class DataListFirst(ComfyNodeABC):
+    """
+    Returns the first element in a list.
+
+    This node takes a list as input and returns the first element of the list.
+    If the list is empty, it returns None.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "list": (IO.ANY, {}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY,)
+    RETURN_NAMES = ("first_element",)
+    CATEGORY = "Basic/Data List"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "get_first_element"
+    INPUT_IS_LIST = True
+
+    def get_first_element(self, **kwargs: list[Any]) -> tuple[Any]:
+        input_list = kwargs.get('list', [])
+        return (input_list[0] if input_list else None,)
+
+
 class DataListGetItem(ComfyNodeABC):
     """
     Retrieves an item at a specified position in a list.
@@ -453,6 +480,33 @@ def insert(self, **kwargs: list[Any]) -> tuple[list[Any]]:
         return (result,)
 
 
+class DataListLast(ComfyNodeABC):
+    """
+    Returns the last element in a list.
+
+    This node takes a list as input and returns the last element of the list.
+    If the list is empty, it returns None.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "list": (IO.ANY, {}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY,)
+    RETURN_NAMES = ("last_element",)
+    CATEGORY = "Basic/Data List"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "get_last_element"
+    INPUT_IS_LIST = True
+
+    def get_last_element(self, **kwargs: list[Any]) -> tuple[Any]:
+        input_list = kwargs.get('list', [])
+        return (input_list[-1] if input_list else None,)
+
+
 class DataListLength(ComfyNodeABC):
     """
     Counts the number of items in a list.
@@ -588,6 +642,38 @@ def pop(self, **kwargs: list[Any]) -> tuple[list[Any], Any]:
             return result, None
 
 
+class DataListPopRandom(ComfyNodeABC):
+    """
+    Removes and returns a random element from a list.
+
+    This node takes a list as input and returns the list with the random element removed
+    and the removed element itself. If the list is empty, it returns None for the element.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "list": (IO.ANY, {}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY, IO.ANY)
+    RETURN_NAMES = ("list", "item")
+    CATEGORY = "Basic/Data List"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "pop_random_element"
+    INPUT_IS_LIST = True
+    OUTPUT_IS_LIST = (True, False)
+
+    def pop_random_element(self, **kwargs: list[Any]) -> tuple[list[Any], Any]:
+        from random import randrange
+        input_list = kwargs.get('list', []).copy()
+        if input_list:
+            random_element = input_list.pop(randrange(len(input_list)))
+            return input_list, random_element
+        return input_list, None
+
+
 class DataListRemove(ComfyNodeABC):
     """
     Removes the first occurrence of a specified value from a list.
@@ -864,13 +950,16 @@ def convert(self, **kwargs: list[Any]) -> tuple[set[Any]]:
     "Basic data handling: DataListExtend": DataListExtend,
     "Basic data handling: DataListFilter": DataListFilter,
     "Basic data handling: DataListFilterSelect": DataListFilterSelect,
+    "Basic data handling: DataListFirst": DataListFirst,
     "Basic data handling: DataListGetItem": DataListGetItem,
     "Basic data handling: DataListIndex": DataListIndex,
     "Basic data handling: DataListInsert": DataListInsert,
+    "Basic data handling: DataListLast": DataListLast,
     "Basic data handling: DataListLength": DataListLength,
     "Basic data handling: DataListMax": DataListMax,
     "Basic data handling: DataListMin": DataListMin,
     "Basic data handling: DataListPop": DataListPop,
+    "Basic data handling: DataListPopRandom": DataListPopRandom,
     "Basic data handling: DataListRemove": DataListRemove,
     "Basic data handling: DataListReverse": DataListReverse,
     "Basic data handling: DataListSetItem": DataListSetItem,
@@ -893,13 +982,16 @@ def convert(self, **kwargs: list[Any]) -> tuple[set[Any]]:
     "Basic data handling: DataListExtend": "extend",
     "Basic data handling: DataListFilter": "filter",
     "Basic data handling: DataListFilterSelect": "filter select",
+    "Basic data handling: DataListFirst": "first",
     "Basic data handling: DataListGetItem": "get item",
     "Basic data handling: DataListIndex": "index",
     "Basic data handling: DataListInsert": "insert",
+    "Basic data handling: DataListLast": "last",
     "Basic data handling: DataListLength": "length",
     "Basic data handling: DataListMax": "max",
     "Basic data handling: DataListMin": "min",
     "Basic data handling: DataListPop": "pop",
+    "Basic data handling: DataListPopRandom": "pop random",
     "Basic data handling: DataListRemove": "remove",
     "Basic data handling: DataListReverse": "reverse",
     "Basic data handling: DataListSetItem": "set item",

src/basic_data_handling/dict_nodes.py

@@ -719,6 +719,42 @@ def popitem(self, input_dict: dict) -> tuple[dict, str, Any, bool]:
             return result, "", None, False
 
 
+class DictPopRandom(ComfyNodeABC):
+    """
+    Removes and returns a random key-value pair from a dictionary.
+
+    This node takes a dictionary as input, removes a random key-value pair,
+    and returns the modified dictionary along with the removed key and value.
+    If the dictionary is empty, it returns empty values.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input_dict": ("DICT", {}),
+            }
+        }
+
+    RETURN_TYPES = ("DICT", IO.STRING, IO.ANY, IO.BOOLEAN)
+    RETURN_NAMES = ("dict", "key", "value", "success")
+    CATEGORY = "Basic/DICT"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "pop_random"
+
+    def pop_random(self, input_dict: dict) -> tuple[dict, str, Any, bool]:
+        import random
+        result = input_dict.copy()
+        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
+
+
 class DictRemove(ComfyNodeABC):
     """
     Removes a key-value pair from a dictionary.
@@ -883,6 +919,7 @@ def values(self, input_dict: dict) -> tuple[list]:
     "Basic data handling: DictMerge": DictMerge,
     "Basic data handling: DictPop": DictPop,
     "Basic data handling: DictPopItem": DictPopItem,
+    "Basic data handling: DictPopRandom": DictPopRandom,
     "Basic data handling: DictRemove": DictRemove,
     "Basic data handling: DictSet": DictSet,
     "Basic data handling: DictSetDefault": DictSetDefault,
@@ -913,7 +950,8 @@ def values(self, input_dict: dict) -> tuple[list]:
     "Basic data handling: DictLength": "length",
     "Basic data handling: DictMerge": "merge",
     "Basic data handling: DictPop": "pop",
-    "Basic data handling: DictPopItem": "popitem",
+    "Basic data handling: DictPopItem": "pop item",
+    "Basic data handling: DictPopRandom": "pop random",
     "Basic data handling: DictRemove": "remove",
     "Basic data handling: DictSet": "set",
     "Basic data handling: DictSetDefault": "setdefault",

src/basic_data_handling/list_nodes.py

@@ -248,6 +248,31 @@ def extend(self, list1: list[Any], list2: list[Any]) -> tuple[list[Any]]:
         return (result,)
 
 
+class ListFirst(ComfyNodeABC):
+    """
+    Returns the first element in a LIST.
+
+    This node takes a LIST as input and returns the first element of the list.
+    If the LIST is empty, it returns None.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "list": ("LIST", {}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY,)
+    RETURN_NAMES = ("first_element",)
+    CATEGORY = "Basic/LIST"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "get_first_element"
+
+    def get_first_element(self, list: list[Any]) -> tuple[Any]:
+        return (list[0] if list else None,)
+
+
 class ListGetItem(ComfyNodeABC):
     """
     Retrieves an item at a specified position in a LIST.
@@ -343,6 +368,31 @@ def insert(self, list: list[Any], index: int, item: Any) -> tuple[list[Any]]:
         return (result,)
 
 
+class ListLast(ComfyNodeABC):
+    """
+    Returns the last element in a LIST.
+
+    This node takes a LIST as input and returns the last element of the list.
+    If the LIST is empty, it returns None.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "list": ("LIST", {}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY,)
+    RETURN_NAMES = ("last_element",)
+    CATEGORY = "Basic/LIST"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "get_last_element"
+
+    def get_last_element(self, list: list[Any]) -> tuple[Any]:
+        return (list[-1] if list else None,)
+
+
 class ListLength(ComfyNodeABC):
     """
     Returns the number of items in a LIST.
@@ -466,6 +516,37 @@ def pop(self, list: list[Any], index: int = -1) -> tuple[list[Any], Any]:
             return result, None
 
 
+class ListPopRandom(ComfyNodeABC):
+    """
+    Removes and returns a random element from a LIST.
+
+    This node takes a LIST as input and returns the LIST with the random element removed
+    and the removed element itself. If the LIST is empty, it returns None for the element.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "list": ("LIST", {}),
+            }
+        }
+
+    RETURN_TYPES = ("LIST", IO.ANY)
+    RETURN_NAMES = ("list", "item")
+    CATEGORY = "Basic/LIST"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "pop_random_element"
+
+    def pop_random_element(self, list: list[Any]) -> tuple[list[Any], Any]:
+        import random
+        result = list.copy()
+        if result:
+            random_index = random.randrange(len(result))
+            random_element = result.pop(random_index)
+            return result, random_element
+        return result, None
+
+
 class ListRemove(ComfyNodeABC):
     """
     Removes the first occurrence of a specified value from a LIST.
@@ -680,13 +761,16 @@ def convert(self, list: list[Any]) -> tuple[set[Any]]:
     "Basic data handling: ListContains": ListContains,
     "Basic data handling: ListCount": ListCount,
     "Basic data handling: ListExtend": ListExtend,
+    "Basic data handling: ListFirst": ListFirst,
     "Basic data handling: ListGetItem": ListGetItem,
     "Basic data handling: ListIndex": ListIndex,
     "Basic data handling: ListInsert": ListInsert,
+    "Basic data handling: ListLast": ListLast,
     "Basic data handling: ListLength": ListLength,
     "Basic data handling: ListMax": ListMax,
     "Basic data handling: ListMin": ListMin,
     "Basic data handling: ListPop": ListPop,
+    "Basic data handling: ListPopRandom": ListPopRandom,
     "Basic data handling: ListRemove": ListRemove,
     "Basic data handling: ListReverse": ListReverse,
     "Basic data handling: ListSetItem": ListSetItem,
@@ -706,13 +790,16 @@ def convert(self, list: list[Any]) -> tuple[set[Any]]:
     "Basic data handling: ListContains": "contains",
     "Basic data handling: ListCount": "count",
     "Basic data handling: ListExtend": "extend",
+    "Basic data handling: ListFirst": "first",
     "Basic data handling: ListGetItem": "get item",
     "Basic data handling: ListIndex": "index",
     "Basic data handling: ListInsert": "insert",
+    "Basic data handling: ListLast": "last",
     "Basic data handling: ListLength": "length",
     "Basic data handling: ListMax": "max",
     "Basic data handling: ListMin": "min",
     "Basic data handling: ListPop": "pop",
+    "Basic data handling: ListPopRandom": "pop random",
     "Basic data handling: ListRemove": "remove",
     "Basic data handling: ListReverse": "reverse",
     "Basic data handling: ListSetItem": "set item",