From 883354b26369d845536b1215ba840e938115cef8 Mon Sep 17 00:00:00 2001
From: Ethnogeny <111099761+050011-code@users.noreply.github.com>
Date: Thu, 4 Jul 2024 09:50:59 +1000
Subject: [PATCH] Blender thumbnail support (#273)

* Update thumb_renderer.py

Included support for rendering blender thumbnails

* Add files via upload

Add functions that get the thumbnail's data

* Update thumb_renderer.py

* Update blender_thumbnailer.py

* Update thumb_renderer.py

* Update thumb_renderer.py

Changed where imports are according to feedback

* Update thumb_renderer.py

Changed blender thumbnail function name to reduce ambiguity

* Update blender_thumbnailer.py

Updated function name

* Update blender_thumbnailer.py

* Update blender_thumbnailer.py

Ruff format

* Update thumb_renderer.py

Ruff format

* Update constants.py

Add .blend1, 2, 3 etc file support

* Update blender_thumbnailer.py

Refactor to follow requested changes

* Update thumb_renderer.py

More refactoring

* Update blender_thumbnailer.py

Ruff format

* Update thumb_renderer.py

Ruff format
---
 tagstudio/src/core/constants.py               |  35 ++++++
 .../src/qt/helpers/blender_thumbnailer.py     | 109 ++++++++++++++++++
 tagstudio/src/qt/widgets/thumb_renderer.py    |  32 +++++
 3 files changed, 176 insertions(+)
 create mode 100644 tagstudio/src/qt/helpers/blender_thumbnailer.py

diff --git a/tagstudio/src/core/constants.py b/tagstudio/src/core/constants.py
index d07489aa..509ae1e8 100644
--- a/tagstudio/src/core/constants.py
+++ b/tagstudio/src/core/constants.py
@@ -105,6 +105,41 @@ ARCHIVE_TYPES: list[str] = [
     ".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"]
 
diff --git a/tagstudio/src/qt/helpers/blender_thumbnailer.py b/tagstudio/src/qt/helpers/blender_thumbnailer.py
new file mode 100644
index 00000000..89f17158
--- /dev/null
+++ b/tagstudio/src/qt/helpers/blender_thumbnailer.py
@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+
+# ##### BEGIN GPL LICENSE BLOCK #####
+#
+#  This program is free software; you can redistribute it and/or
+#  modify it under the terms of the GNU General Public License
+#  as published by the Free Software Foundation; either version 2
+#  of the License, or (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software Foundation,
+#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+#
+# ##### END GPL LICENSE BLOCK #####
+
+# <pep8 compliant>
+
+
+## This file is a modified script that gets the thumbnail data stored in a blend file
+
+
+import struct
+from PIL import (
+    Image,
+    ImageOps,
+)
+import gzip
+import os
+
+
+def blend_extract_thumb(path):
+    REND = b"REND"
+    TEST = b"TEST"
+
+    blendfile = open(path, "rb")
+
+    head = blendfile.read(12)
+
+    if head[0:2] == b"\x1f\x8b":  # gzip magic
+        blendfile.close()
+        blendfile = gzip.GzipFile("", "rb", 0, open(path, "rb"))
+        head = blendfile.read(12)
+
+    if not head.startswith(b"BLENDER"):
+        blendfile.close()
+        return None, 0, 0
+
+    is_64_bit = head[7] == b"-"[0]
+
+    # true for PPC, false for X86
+    is_big_endian = head[8] == b"V"[0]
+
+    # blender pre 2.5 had no thumbs
+    if head[9:11] <= b"24":
+        return None, 0, 0
+
+    sizeof_bhead = 24 if is_64_bit else 20
+    int_endian = ">i" if is_big_endian else "<i"
+    int_endian_pair = int_endian + "i"
+
+    while True:
+        bhead = blendfile.read(sizeof_bhead)
+
+        if len(bhead) < sizeof_bhead:
+            return None, 0, 0
+
+        code = bhead[:4]
+        length = struct.unpack(int_endian, bhead[4:8])[0]  # 4 == sizeof(int)
+
+        if code == REND:
+            blendfile.seek(length, os.SEEK_CUR)
+        else:
+            break
+
+    if code != TEST:
+        return None, 0, 0
+
+    try:
+        x, y = struct.unpack(int_endian_pair, blendfile.read(8))  # 8 == sizeof(int) * 2
+    except struct.error:
+        return None, 0, 0
+
+    length -= 8  # sizeof(int) * 2
+
+    if length != x * y * 4:
+        return None, 0, 0
+
+    image_buffer = blendfile.read(length)
+
+    if len(image_buffer) != length:
+        return None, 0, 0
+
+    return image_buffer, x, y
+
+
+def blend_thumb(file_in):
+    buf, width, height = blend_extract_thumb(file_in)
+    image = Image.frombuffer(
+        "RGBA",
+        (width, height),
+        buf,
+    )
+    image = ImageOps.flip(image)
+    return image
diff --git a/tagstudio/src/qt/widgets/thumb_renderer.py b/tagstudio/src/qt/widgets/thumb_renderer.py
index 1e6a3ad1..9ed612c9 100644
--- a/tagstudio/src/qt/widgets/thumb_renderer.py
+++ b/tagstudio/src/qt/widgets/thumb_renderer.py
@@ -28,8 +28,10 @@ from src.core.constants import (
     VIDEO_TYPES,
     IMAGE_TYPES,
     RAW_IMAGE_TYPES,
+    BLENDER_TYPES,
 )
 from src.core.utils.encoding import detect_char_encoding
+from src.qt.helpers.blender_thumbnailer import blend_thumb
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
@@ -204,6 +206,36 @@ class ThumbRenderer(QObject):
                 # 	img_buf = io.BytesIO()
                 # 	plt.savefig(img_buf, format='png')
                 # 	image = Image.open(img_buf)
+
+                # Blender ===========================================================
+                elif _filepath.suffix.lower() in BLENDER_TYPES:
+                    try:
+                        blend_image = blend_thumb(str(_filepath))
+
+                        bg = Image.new("RGB", blend_image.size, color="#1e1e1e")
+                        bg.paste(blend_image, mask=blend_image.getchannel(3))
+                        image = bg
+
+                    except (
+                        AttributeError,
+                        UnidentifiedImageError,
+                        FileNotFoundError,
+                        TypeError,
+                    ) as e:
+                        if str(e) == "expected string or buffer":
+                            logging.info(
+                                f"[ThumbRenderer]{ERROR} {_filepath.name} Doesn't have thumbnail saved. ({type(e).__name__})"
+                            )
+
+                        else:
+                            logging.info(
+                                f"[ThumbRenderer]{ERROR}: Couldn't render thumbnail for {_filepath.name} ({type(e).__name__})"
+                            )
+
+                        image = ThumbRenderer.thumb_file_default_512.resize(
+                            (adj_size, adj_size), resample=Image.Resampling.BILINEAR
+                        )
+
                 # No Rendered Thumbnail ========================================
                 else:
                     image = ThumbRenderer.thumb_file_default_512.resize(