This commit is contained in:
Christian Byrne 2025-04-11 16:08:19 +00:00 committed by GitHub
commit 93037d363b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 6 deletions

View File

@ -85,6 +85,7 @@ cache_helper = CacheHelper()
extension_mimetypes_cache = {
"webp" : "image",
"fbx" : "model",
}
def map_legacy(folder_name: str) -> str:
@ -140,11 +141,14 @@ def get_directory_by_type(type_name: str) -> str | None:
return get_input_directory()
return None
def filter_files_content_types(files: list[str], content_types: Literal["image", "video", "audio"]) -> list[str]:
def filter_files_content_types(files: list[str], content_types: Literal["image", "video", "audio", "model"]) -> list[str]:
"""
Example:
files = os.listdir(folder_paths.get_input_directory())
filter_files_content_types(files, ["image", "audio", "video"])
videos = filter_files_content_types(files, ["video"])
Note:
- 'model' in MIME context refers to 3D models, not files containing trained weights and parameters
"""
global extension_mimetypes_cache
result = []

View File

@ -1,14 +1,17 @@
import pytest
import os
import tempfile
from folder_paths import filter_files_content_types
from folder_paths import filter_files_content_types, extension_mimetypes_cache
from unittest.mock import patch
@pytest.fixture(scope="module")
def file_extensions():
return {
'image': ['gif', 'heif', 'ico', 'jpeg', 'jpg', 'png', 'pnm', 'ppm', 'svg', 'tiff', 'webp', 'xbm', 'xpm'],
'audio': ['aif', 'aifc', 'aiff', 'au', 'flac', 'm4a', 'mp2', 'mp3', 'ogg', 'snd', 'wav'],
'video': ['avi', 'm2v', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ogv', 'qt', 'webm', 'wmv']
'video': ['avi', 'm2v', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ogv', 'qt', 'webm', 'wmv'],
'model': ['gltf', 'glb', 'obj', 'mtl', 'fbx', 'stl']
}
@ -22,7 +25,18 @@ def mock_dir(file_extensions):
yield directory
def test_categorizes_all_correctly(mock_dir, file_extensions):
@pytest.fixture
def patched_mimetype_cache(file_extensions):
# Mock model file extensions since they may not be in the test-runner system's mimetype cache
new_cache = extension_mimetypes_cache.copy()
for extension in file_extensions["model"]:
new_cache[extension] = "model"
with patch("folder_paths.extension_mimetypes_cache", new_cache):
yield
def test_categorizes_all_correctly(mock_dir, file_extensions, patched_mimetype_cache):
files = os.listdir(mock_dir)
for content_type, extensions in file_extensions.items():
filtered_files = filter_files_content_types(files, [content_type])
@ -30,7 +44,7 @@ def test_categorizes_all_correctly(mock_dir, file_extensions):
assert f"sample_{content_type}.{extension}" in filtered_files
def test_categorizes_all_uniquely(mock_dir, file_extensions):
def test_categorizes_all_uniquely(mock_dir, file_extensions, patched_mimetype_cache):
files = os.listdir(mock_dir)
for content_type, extensions in file_extensions.items():
filtered_files = filter_files_content_types(files, [content_type])