Refator: tasks for less duplicated code

backend/apps/files/tasks.py

@@ -1,7 +1,8 @@
 import os
 import io
-import hashlib
 import subprocess
+from pathlib import Path
+from typing import Optional, Tuple
 
 from django.db import transaction
 
@@ -11,12 +12,91 @@ from PIL import Image as PillowImage
 import blurhash
 
 from .models import PostFileModel
+from utils.hash import compute_file_hash_blake3, compute_md5_hash, compute_blur_hash
 
-from utils.hash import compute_file_hash_blake3, compute_blur_hash
+
+class ThumbnailGenerationError(Exception):
+    """Custom exception for thumbnail generation errors."""
+
+    pass
+
+
+def _setup_output_path(file_hash: str, prefix: str = "thumbnail") -> Tuple[str, str]:
+    """
+    Set up the output directory and generate a unique filename.
+
+    Args:
+        file_hash (str): Hash to use in the filename
+        prefix (str): Prefix for the filename
+
+    Returns:
+        Tuple[str, str]: Output directory path and full file path
+    """
+    output_dir = "/tmp/thumbgen/"
+    os.makedirs(output_dir, exist_ok=True)
+
+    filename = f"{prefix}_{file_hash}.png"
+    filepath = os.path.join(output_dir, filename)
+
+    return output_dir, filepath
+
+
+def _update_file_model(
+    file_model: PostFileModel, thumbnail_path: str, thumbnail_filename: str
+) -> None:
+    """
+    Update the PostFileModel with the new thumbnail and related hashes.
+
+    Args:
+        file_model (PostFileModel): The model to update
+        thumbnail_path (str): Path to the generated thumbnail
+        thumbnail_filename (str): Filename for the saved thumbnail
+    """
+    # Compute the hash for the generated thumbnail
+    thumbnail_hash_blake3 = compute_file_hash_blake3(thumbnail_path)
+
+    # Update the PostFileModel's thumbnail field with the new file
+    with open(thumbnail_path, "rb") as file:
+        file_model.thumbnail.save(thumbnail_filename, file)
+
+    # Set the thumbnail hash
+    file_model.thumbnail_hash_blake3 = thumbnail_hash_blake3
+
+    # Generate and set the blur hash for the thumbnail
+    file_model.thumbnail_blur_hash = compute_blur_hash(thumbnail_path)
+
+    # Save the model
+    file_model.save()
+
+
+def _handle_task_error(e: Exception, file_id: int, process_name: str):
+    """
+    Handle errors in thumbnail generation tasks.
+
+    Args:
+        e (Exception): The exception that occurred
+        file_id (int): ID of the file being processed
+        process_name (str): Name of the process for error reporting
+
+    Raises:
+        Retry: To trigger Celery retry mechanism
+    """
+    error_message = f"Error in {process_name} for file {file_id}: {str(e)}"
+    print(error_message)
+    raise Retry(exc=e)
 
 
 @shared_task(autoretry_for=(Exception,), retry_backoff=True, max_retries=5)
-def generate_blur_hash_PostFile(file_id):
+def generate_blur_hash_PostFile(file_id: int) -> str:
+    """
+    Generate and save a blur hash for an image stored in PostFileModel.
+
+    Args:
+        file_id (int): ID of the PostFileModel instance
+
+    Returns:
+        str: Success message
+    """
     try:
         with transaction.atomic():
             img = PostFileModel.objects.select_for_update().get(id=file_id)
@@ -28,47 +108,56 @@ def generate_blur_hash_PostFile(file_id):
             img.refresh_from_db()
             img.blur_hash = blurhash_string
             img.save()
-        return f"Successfully generated blur hash for file {file_id}"  # Success message
+
+        return f"Successfully generated blur hash for file {file_id}"
+
     except Exception as e:
-        error_message = f"Error generating blur hash for file {file_id}: {e}"
+        _handle_task_error(e, file_id, "blur hash generation")
-        print(error_message)
-        raise Retry(exc=e)  # Retry on exception
-        return error_message  # This ensures the error message is stored in the results
 
 
 @shared_task(autoretry_for=(Exception,), retry_backoff=True, max_retries=5)
-def generate_md5_hash_PostFile(file_id):
+def generate_md5_hash_PostFile(file_id: int) -> str:
+    """
+    Generate and save an MD5 hash for a file stored in PostFileModel.
+
+    Args:
+        file_id (int): ID of the PostFileModel instance
+
+    Returns:
+        str: Success message
+    """
     try:
         with transaction.atomic():
             pstfile = PostFileModel.objects.select_for_update().get(id=file_id)
-            hash_md5 = hashlib.md5()
-            with open(pstfile.file.path, "rb") as f:
-                for chunk in iter(lambda: f.read(4096), b""):
-                    hash_md5.update(chunk)
-            md5_hash = hash_md5.hexdigest()
 
+            # Compute the MD5 hash
+            md5_hash = compute_md5_hash(pstfile.file.path)
+
+            # Save the computed hash
             pstfile.refresh_from_db()
             pstfile.hash_md5 = md5_hash
             pstfile.save()
-        return f"Successfully generated MD5 hash for file {file_id}"  # Success message
+
+        return f"Successfully generated MD5 hash for file {file_id}"
+
     except Exception as e:
-        error_message = f"Error generating MD5 hash for file {file_id}: {e}"
+        _handle_task_error(e, file_id, "MD5 hash generation")
-        print(error_message)
-        raise Retry(exc=e)  # Retry on exception
-        return error_message  # This ensures the error message is stored in the results
 
 
 @shared_task(name="generate_video_thumbnail")
 def generate_video_thumbnail(
-    file_id: int, size: int = 0, timestamp=None, movie_strip: bool = False
+    file_id: int,
-):
+    size: int = 0,
+    timestamp: Optional[float] = None,
+    movie_strip: bool = False,
+) -> str:
     """
     Generate video thumbnails using ffmpegthumbnailer and update the PostFileModel instance.
 
     Args:
         file_id (int): ID of the PostFileModel instance
-        size (int): Desired thumbnail width or height defulats to video size
+        size (int): Desired thumbnail width or height (defaults to video size)
-        timestamp (float): Timestamp(s) in seconds where the thumbnail should be extracted
+        timestamp (float): Timestamp in seconds where the thumbnail should be extracted
         movie_strip (bool): Create a movie strip overlay
 
     Returns:
@@ -84,14 +173,13 @@ def generate_video_thumbnail(
 
             video_path = pstfile.file.path
 
-            # Create output directory if it doesn't exist
+            # Setup output path
-            output_dir = "/tmp/thumbgen/"
+            _, thumbnail_file_path = _setup_output_path(
-            os.makedirs(output_dir, exist_ok=True)
+                pstfile.hash_blake3, "video_thumbnail"
-
+            )
-            thumbnail_filename = f"thumbnail_{pstfile.hash_blake3}.png"
+            thumbnail_filename = Path(thumbnail_file_path).name
-
-            thumbnail_file_path = os.path.join(output_dir, thumbnail_filename)
 
+            # Build command
             cmd = [
                 "ffmpegthumbnailer",
                 "-i",
@@ -110,25 +198,24 @@ def generate_video_thumbnail(
             if timestamp is not None:
                 cmd.extend(["-t", f"{timestamp}"])
 
+            # Execute command
             subprocess.run(cmd, check=True)
 
-            thumbnail_hash_blake3 = compute_file_hash_blake3(thumbnail_file_path)
+            # Update model with new thumbnail
+            _update_file_model(pstfile, thumbnail_file_path, thumbnail_filename)
 
-            # Update the PostFileModel's thumbnail field with the new file
+            # Clean up temporary file
-            with open(thumbnail_file_path, "rb") as file:
+            os.remove(thumbnail_file_path)
-                pstfile.thumbnail.save(thumbnail_filename, file)
 
-            pstfile.thumbnail_hash_blake3 = thumbnail_hash_blake3
+        return f"Video thumbnail generated successfully for file {file_id}"
 
-            pstfile.thumbnail_blur_hash = compute_blur_hash(thumbnail_file_path)
+    except subprocess.CalledProcessError as e:
+        _handle_task_error(e, file_id, "video thumbnail generation")
+    except Exception as e:
+        _handle_task_error(e, file_id, "video thumbnail generation")
 
-            pstfile.save()
 
             os.remove(thumbnail_file_path)
 
-        return f"Thumbnail generated and saved to {thumbnail_file_path}"
 
-    except subprocess.CalledProcessError as e:
-        return f"Error generating thumbnail: {str(e)}"
     except Exception as e:
-        return f"Unexpected error: {str(e)}"