From 20c7e093cb84f96dc451af11485c1bed1f9f0e47 Mon Sep 17 00:00:00 2001 From: daxcay Date: Fri, 5 Jul 2024 10:29:56 +0530 Subject: [PATCH] Bug fixes --- classes/DataSet_PathSelector.py | 63 ++++++++++++++---------- classes/DataSet_TextFilesLoadFromList.py | 2 +- 2 files changed, 38 insertions(+), 27 deletions(-) diff --git a/classes/DataSet_PathSelector.py b/classes/DataSet_PathSelector.py index 72141f0..25dfdfd 100644 --- a/classes/DataSet_PathSelector.py +++ b/classes/DataSet_PathSelector.py @@ -1,18 +1,35 @@ import os -def search_and_select_files(search_in_directory, search_for_extension, select_from_directory, select_extension): - # Search for files with the specified extension in the search_in_directory - search_results = [file for file in os.listdir(search_in_directory) if file.endswith(search_for_extension)] - search_results_names = {os.path.splitext(file)[0] for file in search_results} - # List all files with the specified extension in the select_from_directory - selection_files = [file for file in os.listdir(select_from_directory) if file.endswith(select_extension)] - selection_files_dict = {os.path.splitext(file)[0]: file for file in selection_files} +def normalize_extensions(extensions): + return [ext if ext.startswith('.') else '.' + ext for ext in extensions] + +def search_and_select_files(search_in_directory, search_for_extensions, select_from_directory, select_extensions): + # Normalize extensions + search_for_extensions = normalize_extensions(search_for_extensions.split(',')) + select_extensions = normalize_extensions(select_extensions.split(',')) + + # Search for files with the specified extensions in the search_in_directory + search_results = [ + file for file in os.listdir(search_in_directory) + if any(file.endswith(ext) for ext in search_for_extensions) + ] + search_results_names = {os.path.splitext(file)[0] for file in search_results} + + # List all files with the specified extensions in the select_from_directory + selection_files = [ + file for file in os.listdir(select_from_directory) + if any(file.endswith(ext) for ext in select_extensions) + ] + selection_files_dict = {os.path.splitext(file)[0]: file for file in selection_files} + # Find the matching files based on the names matching_files = [selection_files_dict[name] for name in search_results_names if name in selection_files_dict] + # Prepare the return values matching_files_full_paths = [os.path.join(select_from_directory, file) for file in matching_files] matching_files_with_extensions = matching_files - matching_files_without_extensions = [os.path.splitext(file)[0] for file in matching_files] + matching_files_without_extensions = [os.path.splitext(file)[0] for file in matching_files] + return (matching_files_with_extensions, matching_files_without_extensions, matching_files_full_paths) class DataSet_PathSelector: @@ -22,40 +39,34 @@ class DataSet_PathSelector: @classmethod def INPUT_TYPES(cls): - return { "required": { "search_in_directory": ("STRING", {"multiline": False, "default": ""}), - "search_for_extension": ("STRING", {"multiline": False, "default": ""}), + "search_for_extensions": ("STRING", {"multiline": False, "default": ""}), "select_from_directory": ("STRING", {"multiline": False, "default": ""}), - "select_extension": ("STRING", {"multiline": False, "default": ""}), + "select_extensions": ("STRING", {"multiline": False, "default": ""}), }, } - RETURN_TYPES = ("STRING", "STRING", "INT",) - RETURN_NAMES = ("SelectedNamesWithExtension","SelectedNamesWithoutExtension", "SelectedPaths") + RETURN_TYPES = ("STRING", "STRING", "STRING",) + RETURN_NAMES = ("SelectedNamesWithExtension", "SelectedNamesWithoutExtension", "SelectedPaths") OUTPUT_IS_LIST = (True, True, True) OUTPUT_NODE = True FUNCTION = "make_list" CATEGORY = "🔶DATASET🔶" - def make_list(self, search_in_directory, search_for_extension, select_from_directory, select_extension): - + def make_list(self, search_in_directory, search_for_extensions, select_from_directory, select_extensions): try: - if not os.path.exists(search_in_directory): - print(f"The folder '{search_in_directory}' does not exist.") - return ([], [], 0) + print(f"The folder '{search_in_directory}' does not exist.") + return ([], [], []) if not os.path.exists(select_from_directory): - print(f"The folder '{select_from_directory}' does not exist.") - return ([], [], 0) - - search_for_extension = search_for_extension if search_for_extension.startswith('.') else '.' + search_for_extension - select_extension = select_extension if select_extension.startswith('.') else '.' + select_extension + print(f"The folder '{select_from_directory}' does not exist.") + return ([], [], []) - a,b,c = search_and_select_files(search_in_directory, search_for_extension, select_from_directory, select_extension) - return (a,b,c) + a, b, c = search_and_select_files(search_in_directory, search_for_extensions, select_from_directory, select_extensions) + return (a, b, c) except Exception as e: print(f"An error occurred: {e}") @@ -67,4 +78,4 @@ N_CLASS_MAPPINGS = { N_DISPLAY_NAME_MAPPINGS = { "DataSet_PathSelector": "DataSet_PathSelector", -} \ No newline at end of file +} diff --git a/classes/DataSet_TextFilesLoadFromList.py b/classes/DataSet_TextFilesLoadFromList.py index 21e1c3e..8825be3 100644 --- a/classes/DataSet_TextFilesLoadFromList.py +++ b/classes/DataSet_TextFilesLoadFromList.py @@ -26,7 +26,7 @@ class DataSet_TextFilesLoadFromList: } INPUT_IS_LIST = True - RETURN_TYPES = ("STRING","STRING","STRING",) + RETURN_TYPES = ("STRING","STRING","STRING","STRING") RETURN_NAMES = ("TextFileNames","TextFileNamesWithoutExtension","TextFilePaths","TextFileContents",) OUTPUT_IS_LIST = (True, True, True, True,) FUNCTION = "LoadIT"