82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
import os
|
|
|
|
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]
|
|
|
|
return (matching_files_with_extensions, matching_files_without_extensions, matching_files_full_paths)
|
|
|
|
class DataSet_PathSelector:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"search_in_directory": ("STRING", {"multiline": False, "default": ""}),
|
|
"search_for_extensions": ("STRING", {"multiline": False, "default": ""}),
|
|
"select_from_directory": ("STRING", {"multiline": False, "default": ""}),
|
|
"select_extensions": ("STRING", {"multiline": False, "default": ""}),
|
|
},
|
|
}
|
|
|
|
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_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 ([], [], [])
|
|
|
|
if not os.path.exists(select_from_directory):
|
|
print(f"The folder '{select_from_directory}' does not exist.")
|
|
return ([], [], [])
|
|
|
|
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}")
|
|
return ([], [], [])
|
|
|
|
N_CLASS_MAPPINGS = {
|
|
"DataSet_PathSelector": DataSet_PathSelector,
|
|
}
|
|
|
|
N_DISPLAY_NAME_MAPPINGS = {
|
|
"DataSet_PathSelector": "DataSet_PathSelector",
|
|
}
|