diff --git a/tagstudio/src/core/constants.py b/tagstudio/src/core/constants.py index 2bd1fc379..ef57d7c5d 100644 --- a/tagstudio/src/core/constants.py +++ b/tagstudio/src/core/constants.py @@ -12,156 +12,6 @@ ) FONT_SAMPLE_SIZES: list[int] = [10, 15, 20] -# TODO: Turn this whitelist into a user-configurable blacklist. -IMAGE_TYPES: list[str] = [ - ".png", - ".jpg", - ".jpeg", - ".jpg_large", - ".jpeg_large", - ".jfif", - ".gif", - ".tif", - ".tiff", - ".heic", - ".heif", - ".webp", - ".bmp", - ".svg", - ".avif", - ".apng", - ".jp2", - ".j2k", - ".jpg2", -] -RAW_IMAGE_TYPES: list[str] = [ - ".raw", - ".dng", - ".rw2", - ".nef", - ".arw", - ".crw", - ".cr2", - ".cr3", -] -VIDEO_TYPES: list[str] = [ - ".mp4", - ".webm", - ".mov", - ".hevc", - ".mkv", - ".avi", - ".wmv", - ".flv", - ".gifv", - ".m4p", - ".m4v", - ".3gp", -] -AUDIO_TYPES: list[str] = [ - ".mp3", - ".mp4", - ".mpeg4", - ".m4a", - ".aac", - ".wav", - ".flac", - ".alac", - ".wma", - ".ogg", - ".aiff", -] -DOC_TYPES: list[str] = [ - ".txt", - ".rtf", - ".md", - ".doc", - ".docx", - ".pdf", - ".tex", - ".odt", - ".pages", -] -PLAINTEXT_TYPES: list[str] = [ - ".txt", - ".md", - ".css", - ".html", - ".xml", - ".json", - ".js", - ".ts", - ".ini", - ".htm", - ".csv", - ".php", - ".sh", - ".bat", -] -SPREADSHEET_TYPES: list[str] = [".csv", ".xls", ".xlsx", ".numbers", ".ods"] -PRESENTATION_TYPES: list[str] = [".ppt", ".pptx", ".key", ".odp"] -ARCHIVE_TYPES: list[str] = [ - ".zip", - ".rar", - ".tar", - ".tar", - ".gz", - ".tgz", - ".7z", - ".s7z", -] -BLENDER_TYPES: list[str] = [ - ".blend", - ".blend1", - ".blend2", - ".blend3", - ".blend4", - ".blend5", - ".blend6", - ".blend7", - ".blend8", - ".blend9", - ".blend10", - ".blend11", - ".blend12", - ".blend13", - ".blend14", - ".blend15", - ".blend16", - ".blend17", - ".blend18", - ".blend19", - ".blend20", - ".blend21", - ".blend22", - ".blend23", - ".blend24", - ".blend25", - ".blend26", - ".blend27", - ".blend28", - ".blend29", - ".blend30", - ".blend31", - ".blend32", -] -PROGRAM_TYPES: list[str] = [".exe", ".app"] -SHORTCUT_TYPES: list[str] = [".lnk", ".desktop", ".url"] -FONT_TYPES: list[str] = [".ttf", ".otf", ".woff", ".woff2", ".ttc"] - -ALL_FILE_TYPES: list[str] = ( - IMAGE_TYPES - + VIDEO_TYPES - + AUDIO_TYPES - + DOC_TYPES - + SPREADSHEET_TYPES - + PRESENTATION_TYPES - + ARCHIVE_TYPES - + PROGRAM_TYPES - + SHORTCUT_TYPES - + FONT_TYPES -) - BOX_FIELDS = ["tag_box", "text_box"] TEXT_FIELDS = ["text_line", "text_box"] DATE_FIELDS = ["datetime"] diff --git a/tagstudio/src/core/media_types.py b/tagstudio/src/core/media_types.py new file mode 100644 index 000000000..0bea5e42e --- /dev/null +++ b/tagstudio/src/core/media_types.py @@ -0,0 +1,408 @@ +# Copyright (C) 2024 Travis Abendshien (CyanVoxel). +# Licensed under the GPL-3.0 License. +# Created for TagStudio: https://github.com/CyanVoxel/TagStudio + +import logging +import mimetypes +from enum import Enum +from pathlib import Path + +logging.basicConfig(format="%(message)s", level=logging.INFO) + + +class MediaType(str, Enum): + """Names of media types.""" + + ARCHIVE: str = "archive" + AUDIO: str = "audio" + BLENDER: str = "blender" + DATABASE: str = "database" + DISK_IMAGE: str = "disk_image" + DOCUMENT: str = "document" + FONT: str = "font" + IMAGE_RAW: str = "image_raw" + IMAGE_VECTOR: str = "image_vector" + IMAGE: str = "image" + INSTALLER: str = "installer" + MATERIAL: str = "material" + MODEL: str = "model" + PACKAGE: str = "package" + PHOTOSHOP: str = "photoshop" + PLAINTEXT: str = "plaintext" + PRESENTATION: str = "presentation" + PROGRAM: str = "program" + SHORTCUT: str = "shortcut" + SPREADSHEET: str = "spreadsheet" + TEXT: str = "text" + VIDEO: str = "video" + + +class MediaCategory: + """An object representing a category of media. Includes a MediaType identifier, + extensions set, and IANA status flag. + + Args: + media_type (MediaType): The MediaType Enum representing this category. + + extensions (set[str]): The set of file extensions associated with this category. + Includes leading ".", all lowercase, and does not need to be unique to this category. + + is_iana (bool): Represents whether or not this is an IANA registered category. + """ + + def __init__( + self, + media_type: MediaType, + extensions: set[str], + is_iana: bool = False, + ) -> None: + self.media_type: MediaType = media_type + self.extensions: set[str] = extensions + self.is_iana: bool = is_iana + + +class MediaCategories: + """Contains pre-made MediaCategory objects as well as methods to interact with them.""" + + # These sets are used either individually or together to form the final sets + # for the MediaCategory(s). + # These sets may be combined and are NOT 1:1 with the final categories. + _ARCHIVE_SET: set[str] = { + ".7z", + ".gz", + ".rar", + ".s7z", + ".tar", + ".tgz", + ".zip", + } + _AUDIO_SET: set[str] = { + ".aac", + ".aif", + ".aiff", + ".alac", + ".flac", + ".m4a", + ".m4p", + ".mp3", + ".mpeg4", + ".ogg", + ".wav", + ".wma", + } + _BLENDER_SET: set[str] = { + ".blen_tc", + ".blend", + ".blend1", + ".blend10", + ".blend11", + ".blend12", + ".blend13", + ".blend14", + ".blend15", + ".blend16", + ".blend17", + ".blend18", + ".blend19", + ".blend2", + ".blend20", + ".blend21", + ".blend22", + ".blend23", + ".blend24", + ".blend25", + ".blend26", + ".blend27", + ".blend28", + ".blend29", + ".blend3", + ".blend30", + ".blend31", + ".blend32", + ".blend4", + ".blend5", + ".blend6", + ".blend7", + ".blend8", + ".blend9", + } + _DATABASE_SET: set[str] = { + ".accdb", + ".mdb", + ".sqlite", + } + _DISK_IMAGE_SET: set[str] = {".bios", ".dmg", ".iso"} + _DOCUMENT_SET: set[str] = { + ".doc", + ".docm", + ".docx", + ".dot", + ".dotm", + ".dotx", + ".odt", + ".pages", + ".pdf", + ".rtf", + ".tex", + ".wpd", + ".wps", + } + _FONT_SET: set[str] = { + ".fon", + ".otf", + ".ttc", + ".ttf", + ".woff", + ".woff2", + } + _IMAGE_RAW_SET: set[str] = { + ".arw", + ".cr2", + ".cr3", + ".crw", + ".dng", + ".nef", + ".raw", + ".rw2", + } + _IMAGE_VECTOR_SET: set[str] = {".svg"} + _IMAGE_SET: set[str] = { + ".apng", + ".avif", + ".bmp", + ".exr", + ".gif", + ".heic", + ".heif", + ".j2k", + ".jfif", + ".jp2", + ".jpeg_large", + ".jpeg", + ".jpg_large", + ".jpg", + ".jpg2", + ".png", + ".psb", + ".psd", + ".tif", + ".tiff", + ".webp", + } + _INSTALLER_SET: set[str] = {".appx", ".msi", ".msix"} + _MATERIAL_SET: set[str] = {".mtl"} + _MODEL_SET: set[str] = {".3ds", ".fbx", ".obj", ".stl"} + _PACKAGE_SET: set[str] = {".pkg"} + _PHOTOSHOP_SET: set[str] = { + ".pdd", + ".psb", + ".psd", + } + _PLAINTEXT_SET: set[str] = { + ".bat", + ".css", + ".csv", + ".htm", + ".html", + ".ini", + ".js", + ".json", + ".jsonc", + ".md", + ".php", + ".plist", + ".prefs", + ".sh", + ".ts", + ".txt", + ".xml", + } + _PRESENTATION_SET: set[str] = { + ".ppt", + ".pptx", + ".key", + ".odp", + } + _PROGRAM_SET: set[str] = {".app", ".exe"} + _SHORTCUT_SET: set[str] = {".desktop", ".lnk", ".url"} + _SPREADSHEET_SET: set[str] = { + ".csv", + ".numbers", + ".ods" ".xls", + ".xlsx", + } + _VIDEO_SET: set[str] = { + ".3gp", + ".avi", + ".flv", + ".gifv", + ".hevc", + ".m4p", + ".m4v", + ".mkv", + ".mov", + ".mp4", + ".webm", + ".wmv", + } + + ARCHIVE_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.ARCHIVE, + extensions=_ARCHIVE_SET, + is_iana=False, + ) + AUDIO_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.AUDIO, + extensions=_AUDIO_SET, + is_iana=True, + ) + BLENDER_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.BLENDER, + extensions=_BLENDER_SET, + is_iana=False, + ) + DATABASE_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.DATABASE, + extensions=_DATABASE_SET, + is_iana=False, + ) + DISK_IMAGE_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.DISK_IMAGE, + extensions=_DISK_IMAGE_SET, + is_iana=False, + ) + DOCUMENT_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.DOCUMENT, + extensions=_DOCUMENT_SET, + is_iana=False, + ) + FONT_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.FONT, + extensions=_FONT_SET, + is_iana=True, + ) + IMAGE_RAW_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.IMAGE_RAW, + extensions=_IMAGE_RAW_SET, + is_iana=False, + ) + IMAGE_VECTOR_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.IMAGE_VECTOR, + extensions=_IMAGE_VECTOR_SET, + is_iana=False, + ) + IMAGE_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.IMAGE, + extensions=_IMAGE_SET | _IMAGE_RAW_SET | _IMAGE_VECTOR_SET, + is_iana=True, + ) + INSTALLER_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.INSTALLER, + extensions=_INSTALLER_SET, + is_iana=False, + ) + MATERIAL_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.MATERIAL, + extensions=_MATERIAL_SET, + is_iana=False, + ) + MODEL_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.MODEL, + extensions=_MODEL_SET, + is_iana=True, + ) + PACKAGE_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.PACKAGE, + extensions=_PACKAGE_SET, + is_iana=False, + ) + PHOTOSHOP_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.PHOTOSHOP, + extensions=_PHOTOSHOP_SET, + is_iana=False, + ) + PLAINTEXT_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.PLAINTEXT, + extensions=_PLAINTEXT_SET, + is_iana=False, + ) + PRESENTATION_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.PRESENTATION, + extensions=_PRESENTATION_SET, + is_iana=False, + ) + PROGRAM_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.PROGRAM, + extensions=_PROGRAM_SET, + is_iana=False, + ) + SHORTCUT_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.SHORTCUT, + extensions=_SHORTCUT_SET, + is_iana=False, + ) + SPREADSHEET_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.SPREADSHEET, + extensions=_SPREADSHEET_SET, + is_iana=False, + ) + TEXT_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.TEXT, + extensions=_DOCUMENT_SET | _PLAINTEXT_SET, + is_iana=True, + ) + VIDEO_TYPES: MediaCategory = MediaCategory( + media_type=MediaType.VIDEO, + extensions=_VIDEO_SET, + is_iana=True, + ) + + ALL_CATEGORIES: list[MediaCategory] = [ + ARCHIVE_TYPES, + AUDIO_TYPES, + BLENDER_TYPES, + DATABASE_TYPES, + DISK_IMAGE_TYPES, + DOCUMENT_TYPES, + FONT_TYPES, + IMAGE_RAW_TYPES, + IMAGE_TYPES, + IMAGE_VECTOR_TYPES, + INSTALLER_TYPES, + MATERIAL_TYPES, + MODEL_TYPES, + PACKAGE_TYPES, + PHOTOSHOP_TYPES, + PLAINTEXT_TYPES, + PRESENTATION_TYPES, + PROGRAM_TYPES, + SHORTCUT_TYPES, + SPREADSHEET_TYPES, + TEXT_TYPES, + VIDEO_TYPES, + ] + + @staticmethod + def get_types(ext: str, mime_fallback: bool = False) -> set[MediaType]: + """Returns a set of MediaTypes given a file extension. + + Args: + ext (str): File extension with a leading "." and in all lowercase. + mime_fallback (bool): Flag to guess MIME type if no set matches are made. + """ + types: set[MediaType] = set() + mime_guess: bool = False + + for cat in MediaCategories.ALL_CATEGORIES: + if ext in cat.extensions: + types.add(cat.media_type) + elif mime_fallback and cat.is_iana: + type: str = mimetypes.guess_type(Path("x" + ext), strict=False)[0] + if type and type.startswith(cat.media_type.value): + types.add(cat.media_type) + mime_guess = True + + # logging.info( + # f"({ext}) Media Categories Found: {[x.value for x in types]}{' (MIME)' if mime_guess else ''}" + # ) + return types diff --git a/tagstudio/src/qt/widgets/collage_icon.py b/tagstudio/src/qt/widgets/collage_icon.py index b9234d7d2..ddfda5c06 100644 --- a/tagstudio/src/qt/widgets/collage_icon.py +++ b/tagstudio/src/qt/widgets/collage_icon.py @@ -24,7 +24,7 @@ ) from src.core.library import Library -from src.core.constants import DOC_TYPES, VIDEO_TYPES, IMAGE_TYPES +from src.core.media_types import MediaCategories, MediaType ERROR = f"[ERROR]" @@ -93,7 +93,8 @@ def render( ) # sys.stdout.write(f'\r{INFO} Combining [{i+1}/{len(self.lib.entries)}]: {self.get_file_color(file_type)}{entry.path}{os.sep}{entry.filename}{RESET}') # sys.stdout.flush() - if filepath.suffix.lower() in IMAGE_TYPES: + ext: str = filepath.suffix.lower() + if MediaType.IMAGE in MediaCategories.get_types(ext): try: with Image.open( str(self.lib.library_dir / entry.path / entry.filename) @@ -111,7 +112,7 @@ def render( self.rendered.emit(pic) except DecompressionBombError as e: logging.info(f"[ERROR] One of the images was too big ({e})") - elif filepath.suffix.lower() in VIDEO_TYPES: + elif MediaType.VIDEO in MediaCategories.get_types(ext): video = cv2.VideoCapture(str(filepath)) video.set( cv2.CAP_PROP_POS_FRAMES, @@ -167,14 +168,16 @@ def render( self.done.emit() # logging.info('Done!') + # NOTE: Depreciated def get_file_color(self, ext: str): - if ext.lower().replace(".", "", 1) == "gif": + _ext = ext.lower().replace(".", "", 1) + if _ext == "gif": return "\033[93m" - if ext.lower().replace(".", "", 1) in IMAGE_TYPES: + elif MediaType.IMAGE in MediaCategories.get_types(_ext): return "\033[37m" - elif ext.lower().replace(".", "", 1) in VIDEO_TYPES: + elif MediaType.VIDEO in MediaCategories.get_types(_ext): return "\033[96m" - elif ext.lower().replace(".", "", 1) in DOC_TYPES: + elif MediaType.DOCUMENT in MediaCategories.get_types(_ext): return "\033[92m" else: return "\033[97m" diff --git a/tagstudio/src/qt/widgets/item_thumb.py b/tagstudio/src/qt/widgets/item_thumb.py index 0adcb644e..fdeb6739e 100644 --- a/tagstudio/src/qt/widgets/item_thumb.py +++ b/tagstudio/src/qt/widgets/item_thumb.py @@ -24,12 +24,10 @@ from src.core.enums import FieldID from src.core.library import ItemType, Library, Entry from src.core.constants import ( - AUDIO_TYPES, - VIDEO_TYPES, - IMAGE_TYPES, TAG_FAVORITE, TAG_ARCHIVED, ) +from src.core.media_types import MediaCategories, MediaType from src.qt.flowlayout import FlowWidget from src.qt.helpers.file_opener import FileOpenerHelper from src.qt.widgets.thumb_renderer import ThumbRenderer @@ -358,10 +356,24 @@ def set_mode(self, mode: Optional[ItemType]) -> None: def set_extension(self, ext: str) -> None: if ext and ext.startswith(".") is False: ext = "." + ext - if ext and ext not in IMAGE_TYPES or ext in [".gif", ".apng"]: + if ( + ext + and (MediaType.IMAGE not in MediaCategories.get_types(ext)) + or (MediaType.IMAGE_RAW in MediaCategories.get_types(ext)) + or (MediaType.IMAGE_VECTOR in MediaCategories.get_types(ext)) + or (MediaType.PHOTOSHOP in MediaCategories.get_types(ext)) + or ext + in [ + ".apng", + ".exr", + ".gif", + ] + ): self.ext_badge.setHidden(False) self.ext_badge.setText(ext.upper()[1:]) - if ext in VIDEO_TYPES + AUDIO_TYPES: + if (MediaType.VIDEO in MediaCategories.get_types(ext)) or ( + MediaType.AUDIO in MediaCategories.get_types(ext) + ): self.count_badge.setHidden(False) else: if self.mode == ItemType.ENTRY: diff --git a/tagstudio/src/qt/widgets/preview_panel.py b/tagstudio/src/qt/widgets/preview_panel.py index 6892c7f5e..537ee09f4 100644 --- a/tagstudio/src/qt/widgets/preview_panel.py +++ b/tagstudio/src/qt/widgets/preview_panel.py @@ -31,12 +31,9 @@ from src.core.enums import SettingItems, Theme from src.core.library import Entry, ItemType, Library from src.core.constants import ( - VIDEO_TYPES, - IMAGE_TYPES, - RAW_IMAGE_TYPES, TS_FOLDER_NAME, - FONT_TYPES, ) +from src.core.media_types import MediaCategories, MediaType from src.qt.helpers.file_opener import FileOpenerLabel, FileOpenerHelper, open_file from src.qt.modals.add_field import AddFieldModal from src.qt.widgets.thumb_renderer import ThumbRenderer @@ -520,12 +517,23 @@ def update_widgets(self): self.opener.open_explorer ) - # TODO: Do this somewhere else, this is just here temporarily. + # TODO: Do this all somewhere else, this is just here temporarily. + ext: str = filepath.suffix.lower() try: image = None - if filepath.suffix.lower() in IMAGE_TYPES: + if ( + (MediaType.IMAGE in MediaCategories.get_types(ext)) + and ( + MediaType.IMAGE_RAW + not in MediaCategories.get_types(ext) + ) + and ( + MediaType.IMAGE_VECTOR + not in MediaCategories.get_types(ext) + ) + ): image = Image.open(str(filepath)) - elif filepath.suffix.lower() in RAW_IMAGE_TYPES: + elif MediaType.IMAGE_RAW in MediaCategories.get_types(ext): try: with rawpy.imread(str(filepath)) as raw: rgb = raw.postprocess() @@ -537,7 +545,7 @@ def update_widgets(self): rawpy._rawpy.LibRawFileUnsupportedError, ): pass - elif filepath.suffix.lower() in VIDEO_TYPES: + elif MediaType.VIDEO in MediaCategories.get_types(ext): video = cv2.VideoCapture(str(filepath)) if video.get(cv2.CAP_PROP_FRAME_COUNT) <= 0: raise cv2.error("File is invalid or has 0 frames") @@ -559,33 +567,47 @@ def update_widgets(self): self.preview_vid.show() # Stats for specific file types are displayed here. - if image and filepath.suffix.lower() in ( - IMAGE_TYPES + VIDEO_TYPES + RAW_IMAGE_TYPES - ): - self.dimensions_label.setText( - f"{filepath.suffix.upper()[1:]} • {format_size(filepath.stat().st_size)}\n{image.width} x {image.height} px" + if image and ( + (MediaType.IMAGE in MediaCategories.get_types(ext)) + or (MediaType.VIDEO in MediaCategories.get_types(ext, True)) + or ( + MediaType.IMAGE_RAW + in MediaCategories.get_types(ext, True) ) - elif filepath.suffix.lower() in FONT_TYPES: - font = ImageFont.truetype(filepath) + ): self.dimensions_label.setText( - f"{filepath.suffix.upper()[1:]} • {format_size(filepath.stat().st_size)}\n{font.getname()[0]} ({font.getname()[1]}) " + f"{ext.upper()[1:]} • {format_size(filepath.stat().st_size)}\n{image.width} x {image.height} px" ) + elif MediaType.FONT in MediaCategories.get_types(ext, True): + try: + font = ImageFont.truetype(filepath) + self.dimensions_label.setText( + f"{ext.upper()[1:]} • {format_size(filepath.stat().st_size)}\n{font.getname()[0]} ({font.getname()[1]}) " + ) + except OSError: + self.dimensions_label.setText( + f"{ext.upper()[1:]} • {format_size(filepath.stat().st_size)}" + ) + logging.info( + f"[PreviewPanel][ERROR] Couldn't read font file: {filepath}" + ) else: + self.dimensions_label.setText(f"{ext.upper()[1:]}") self.dimensions_label.setText( - f"{filepath.suffix.upper()[1:]} • {format_size(filepath.stat().st_size)}" + f"{ext.upper()[1:]} • {format_size(filepath.stat().st_size)}" ) if not filepath.is_file(): raise FileNotFoundError except FileNotFoundError as e: - self.dimensions_label.setText(f"{filepath.suffix.upper()[1:]}") + self.dimensions_label.setText(f"{ext.upper()[1:]}") logging.info( f"[PreviewPanel][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})" ) except (FileNotFoundError, cv2.error) as e: - self.dimensions_label.setText(f"{filepath.suffix.upper()}") + self.dimensions_label.setText(f"{ext.upper()}") logging.info( f"[PreviewPanel][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})" ) @@ -594,7 +616,7 @@ def update_widgets(self): DecompressionBombError, ) as e: self.dimensions_label.setText( - f"{filepath.suffix.upper()[1:]} • {format_size(filepath.stat().st_size)}" + f"{ext.upper()[1:]} • {format_size(filepath.stat().st_size)}" ) logging.info( f"[PreviewPanel][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})" diff --git a/tagstudio/src/qt/widgets/thumb_renderer.py b/tagstudio/src/qt/widgets/thumb_renderer.py index 47421b4f3..4b37b0c32 100644 --- a/tagstudio/src/qt/widgets/thumb_renderer.py +++ b/tagstudio/src/qt/widgets/thumb_renderer.py @@ -24,16 +24,8 @@ from PySide6.QtGui import QPixmap from src.qt.helpers.gradient import four_corner_gradient_background from src.qt.helpers.text_wrapper import wrap_full_text -from src.core.constants import ( - PLAINTEXT_TYPES, - FONT_TYPES, - VIDEO_TYPES, - IMAGE_TYPES, - RAW_IMAGE_TYPES, - FONT_SAMPLE_TEXT, - FONT_SAMPLE_SIZES, - BLENDER_TYPES, -) +from src.core.constants import FONT_SAMPLE_SIZES, FONT_SAMPLE_TEXT +from src.core.media_types import MediaType, MediaCategories from src.core.utils.encoding import detect_char_encoding from src.qt.helpers.blender_thumbnailer import blend_thumb @@ -126,47 +118,51 @@ def render( self.updated_ratio.emit(1) elif _filepath: try: + ext: str = _filepath.suffix.lower() # Images ======================================================= - if _filepath.suffix.lower() in IMAGE_TYPES: - try: - image = Image.open(_filepath) - if image.mode != "RGB" and image.mode != "RGBA": - image = image.convert(mode="RGBA") - if image.mode == "RGBA": - new_bg = Image.new("RGB", image.size, color="#1e1e1e") - new_bg.paste(image, mask=image.getchannel(3)) - image = new_bg - - image = ImageOps.exif_transpose(image) - except DecompressionBombError as e: - logging.info( - f"[ThumbRenderer]{WARNING} Couldn't Render thumbnail for {_filepath.name} ({type(e).__name__})" - ) + if MediaType.IMAGE in MediaCategories.get_types(ext, True): + # Raw Images ----------------------------------------------- + if MediaType.IMAGE_RAW in MediaCategories.get_types(ext, True): + try: + with rawpy.imread(str(_filepath)) as raw: + rgb = raw.postprocess() + image = Image.frombytes( + "RGB", + (rgb.shape[1], rgb.shape[0]), + rgb, + decoder_name="raw", + ) + except DecompressionBombError as e: + logging.info( + f"[ThumbRenderer]{WARNING} Couldn't Render thumbnail for {_filepath.name} ({type(e).__name__})" + ) + except ( + rawpy._rawpy.LibRawIOError, + rawpy._rawpy.LibRawFileUnsupportedError, + ) as e: + logging.info( + f"[ThumbRenderer]{ERROR} Couldn't Render thumbnail for raw image {_filepath.name} ({type(e).__name__})" + ) - elif _filepath.suffix.lower() in RAW_IMAGE_TYPES: - try: - with rawpy.imread(str(_filepath)) as raw: - rgb = raw.postprocess() - image = Image.frombytes( - "RGB", - (rgb.shape[1], rgb.shape[0]), - rgb, - decoder_name="raw", + # Normal Images -------------------------------------------- + else: + try: + image = Image.open(_filepath) + if image.mode != "RGB" and image.mode != "RGBA": + image = image.convert(mode="RGBA") + if image.mode == "RGBA": + new_bg = Image.new("RGB", image.size, color="#1e1e1e") + new_bg.paste(image, mask=image.getchannel(3)) + image = new_bg + + image = ImageOps.exif_transpose(image) + except DecompressionBombError as e: + logging.info( + f"[ThumbRenderer]{WARNING} Couldn't Render thumbnail for {_filepath.name} ({type(e).__name__})" ) - except DecompressionBombError as e: - logging.info( - f"[ThumbRenderer]{WARNING} Couldn't Render thumbnail for {_filepath.name} ({type(e).__name__})" - ) - except ( - rawpy._rawpy.LibRawIOError, - rawpy._rawpy.LibRawFileUnsupportedError, - ) as e: - logging.info( - f"[ThumbRenderer]{ERROR} Couldn't Render thumbnail for raw image {_filepath.name} ({type(e).__name__})" - ) # Videos ======================================================= - elif _filepath.suffix.lower() in VIDEO_TYPES: + elif MediaType.VIDEO in MediaCategories.get_types(ext, True): video = cv2.VideoCapture(str(_filepath)) frame_count = video.get(cv2.CAP_PROP_FRAME_COUNT) if frame_count <= 0: @@ -183,7 +179,7 @@ def render( image = Image.fromarray(frame) # Plain Text =================================================== - elif _filepath.suffix.lower() in PLAINTEXT_TYPES: + elif MediaType.PLAINTEXT in MediaCategories.get_types(ext): encoding = detect_char_encoding(_filepath) with open(_filepath, "r", encoding=encoding) as text_file: text = text_file.read(256) @@ -192,7 +188,7 @@ def render( draw.text((16, 16), text, fill=(255, 255, 255)) image = bg # Fonts ======================================================== - elif _filepath.suffix.lower() in FONT_TYPES: + elif MediaType.FONT in MediaCategories.get_types(ext, True): # Scale the sample font sizes to the preview image # resolution,assuming the sizes are tuned for 256px. scaled_sizes: list[int] = [ @@ -245,7 +241,7 @@ def render( # image = Image.open(img_buf) # Blender =========================================================== - elif _filepath.suffix.lower() in BLENDER_TYPES: + elif MediaType.BLENDER in MediaCategories.get_types(ext): try: blend_image = blend_thumb(str(_filepath)) @@ -335,6 +331,7 @@ def render( cv2.error, DecompressionBombError, UnicodeDecodeError, + OSError, ) as e: if e is not UnicodeDecodeError: logging.info(