318 lines
12 KiB
Python
318 lines
12 KiB
Python
import os
|
|
import uuid
|
|
from werkzeug.utils import secure_filename
|
|
from flask import current_app
|
|
from PIL import Image
|
|
from typing import Tuple, Optional
|
|
|
|
|
|
class FileHandler:
|
|
"""Service for handling file uploads and storage."""
|
|
|
|
@staticmethod
|
|
def allowed_file(filename: str) -> bool:
|
|
"""Check if file extension is allowed."""
|
|
return '.' in filename and \
|
|
filename.rsplit('.', 1)[1].lower() in current_app.config['ALLOWED_EXTENSIONS']
|
|
|
|
@staticmethod
|
|
def get_file_type(filename: str) -> str:
|
|
"""Determine file type based on extension."""
|
|
ext = filename.rsplit('.', 1)[1].lower() if '.' in filename else ''
|
|
return 'PDF' if ext == 'pdf' else 'IMAGE'
|
|
|
|
@staticmethod
|
|
def generate_unique_filename(original_filename: str) -> str:
|
|
"""Generate a unique filename while preserving the extension."""
|
|
ext = secure_filename(original_filename).rsplit('.', 1)[1].lower()
|
|
unique_name = f"{uuid.uuid4().hex}.{ext}"
|
|
return unique_name
|
|
|
|
@staticmethod
|
|
def save_cover_image(file, set_number: str) -> Tuple[str, int]:
|
|
"""
|
|
Save cover image for a set.
|
|
|
|
Args:
|
|
file: FileStorage object from request
|
|
set_number: Set number for organizing files
|
|
|
|
Returns:
|
|
Tuple of (relative_path, file_size_bytes)
|
|
"""
|
|
# Create covers directory if it doesn't exist
|
|
covers_dir = os.path.join(current_app.config['UPLOAD_FOLDER'], 'covers')
|
|
os.makedirs(covers_dir, exist_ok=True)
|
|
|
|
# Generate unique filename
|
|
filename = FileHandler.generate_unique_filename(file.filename)
|
|
|
|
# Create set-specific subdirectory
|
|
set_dir = os.path.join(covers_dir, secure_filename(set_number))
|
|
os.makedirs(set_dir, exist_ok=True)
|
|
|
|
# Save file
|
|
file_path = os.path.join(set_dir, filename)
|
|
file.save(file_path)
|
|
|
|
# Get file size
|
|
file_size = os.path.getsize(file_path)
|
|
|
|
# Create thumbnail and optimize
|
|
try:
|
|
img = Image.open(file_path)
|
|
|
|
# Convert RGBA to RGB if necessary
|
|
if img.mode == 'RGBA':
|
|
background = Image.new('RGB', img.size, (255, 255, 255))
|
|
background.paste(img, mask=img.split()[3])
|
|
img = background
|
|
|
|
# Resize if too large (max 800px on longest side)
|
|
max_size = 800
|
|
if max(img.size) > max_size:
|
|
ratio = max_size / max(img.size)
|
|
new_size = tuple(int(dim * ratio) for dim in img.size)
|
|
img = img.resize(new_size, Image.Resampling.LANCZOS)
|
|
img.save(file_path, optimize=True, quality=85)
|
|
file_size = os.path.getsize(file_path)
|
|
|
|
except Exception as e:
|
|
current_app.logger.error(f"Error processing cover image: {e}")
|
|
|
|
# Return relative path from uploads folder
|
|
relative_path = os.path.join('covers', secure_filename(set_number), filename)
|
|
return relative_path.replace('\\', '/'), file_size
|
|
|
|
@staticmethod
|
|
def save_file(file, set_number: str, file_type: str) -> Tuple[str, int, Optional[str]]:
|
|
"""
|
|
Save uploaded file to appropriate directory.
|
|
|
|
Args:
|
|
file: FileStorage object from request
|
|
set_number: LEGO set number for organizing files
|
|
file_type: 'PDF' or 'IMAGE'
|
|
|
|
Returns:
|
|
Tuple of (relative_path, file_size, thumbnail_path)
|
|
"""
|
|
# Generate unique filename
|
|
unique_filename = FileHandler.generate_unique_filename(file.filename)
|
|
|
|
# Determine subdirectory
|
|
subdir = 'pdfs' if file_type == 'PDF' else 'images'
|
|
|
|
# Create set-specific directory
|
|
set_dir = os.path.join(
|
|
current_app.config['UPLOAD_FOLDER'],
|
|
subdir,
|
|
secure_filename(set_number)
|
|
)
|
|
os.makedirs(set_dir, exist_ok=True)
|
|
|
|
# Full file path
|
|
file_path = os.path.join(set_dir, unique_filename)
|
|
|
|
# Save the file
|
|
file.save(file_path)
|
|
|
|
# Get file size
|
|
file_size = os.path.getsize(file_path)
|
|
|
|
# Generate thumbnail
|
|
thumbnail_rel_path = None
|
|
if file_type == 'IMAGE':
|
|
thumb_path = FileHandler.create_thumbnail(file_path)
|
|
if thumb_path:
|
|
# Get relative path for thumbnail
|
|
thumbnail_rel_path = os.path.join(subdir, secure_filename(set_number),
|
|
os.path.basename(thumb_path))
|
|
thumbnail_rel_path = thumbnail_rel_path.replace('\\', '/')
|
|
elif file_type == 'PDF':
|
|
thumb_path = FileHandler.create_pdf_thumbnail(file_path, set_number)
|
|
if thumb_path:
|
|
thumbnail_rel_path = thumb_path.replace('\\', '/')
|
|
|
|
# Return relative path for database storage
|
|
relative_path = os.path.join(subdir, secure_filename(set_number), unique_filename)
|
|
|
|
return relative_path.replace('\\', '/'), file_size, thumbnail_rel_path
|
|
|
|
@staticmethod
|
|
def create_thumbnail(image_path: str, size: Tuple[int, int] = (300, 300)) -> Optional[str]:
|
|
"""
|
|
Create a thumbnail for an image.
|
|
|
|
Args:
|
|
image_path: Path to the original image
|
|
size: Thumbnail size (width, height)
|
|
|
|
Returns:
|
|
Path to thumbnail or None if failed
|
|
"""
|
|
try:
|
|
img = Image.open(image_path)
|
|
img.thumbnail(size, Image.Resampling.LANCZOS)
|
|
|
|
# Generate thumbnail filename
|
|
base, ext = os.path.splitext(image_path)
|
|
thumb_path = f"{base}_thumb{ext}"
|
|
|
|
img.save(thumb_path, optimize=True)
|
|
return thumb_path
|
|
|
|
except Exception as e:
|
|
current_app.logger.error(f'Thumbnail creation failed: {str(e)}')
|
|
return None
|
|
|
|
@staticmethod
|
|
def create_pdf_thumbnail(pdf_path: str, set_number: str, size: Tuple[int, int] = (300, 400)) -> Optional[str]:
|
|
"""
|
|
Create a thumbnail from the first page of a PDF.
|
|
|
|
Args:
|
|
pdf_path: Path to the PDF file
|
|
set_number: Set number for organizing thumbnails
|
|
size: Thumbnail size (width, height)
|
|
|
|
Returns:
|
|
Relative path to thumbnail or None if failed
|
|
"""
|
|
try:
|
|
# Try using PyMuPDF (fitz) first
|
|
try:
|
|
import fitz # PyMuPDF
|
|
|
|
# Open PDF
|
|
doc = fitz.open(pdf_path)
|
|
|
|
# Get first page
|
|
page = doc[0]
|
|
|
|
# Render page to pixmap (image)
|
|
# Use matrix for scaling to desired size
|
|
zoom = 2 # Higher zoom for better quality
|
|
mat = fitz.Matrix(zoom, zoom)
|
|
pix = page.get_pixmap(matrix=mat)
|
|
|
|
# Create thumbnails directory
|
|
thumbnails_dir = os.path.join(current_app.config['UPLOAD_FOLDER'], 'thumbnails', secure_filename(set_number))
|
|
os.makedirs(thumbnails_dir, exist_ok=True)
|
|
|
|
# Generate thumbnail filename
|
|
pdf_filename = os.path.basename(pdf_path)
|
|
thumb_filename = os.path.splitext(pdf_filename)[0] + '_thumb.png'
|
|
thumb_path = os.path.join(thumbnails_dir, thumb_filename)
|
|
|
|
# Save as PNG
|
|
pix.save(thumb_path)
|
|
|
|
# Close document
|
|
doc.close()
|
|
|
|
# Resize to target size using PIL
|
|
img = Image.open(thumb_path)
|
|
img.thumbnail(size, Image.Resampling.LANCZOS)
|
|
img.save(thumb_path, optimize=True)
|
|
|
|
# Return relative path
|
|
rel_path = os.path.join('thumbnails', secure_filename(set_number), thumb_filename)
|
|
return rel_path
|
|
|
|
except ImportError:
|
|
# Try pdf2image as fallback
|
|
try:
|
|
from pdf2image import convert_from_path
|
|
|
|
# Convert first page only
|
|
images = convert_from_path(pdf_path, first_page=1, last_page=1, dpi=150)
|
|
|
|
if images:
|
|
# Create thumbnails directory
|
|
thumbnails_dir = os.path.join(current_app.config['UPLOAD_FOLDER'], 'thumbnails', secure_filename(set_number))
|
|
os.makedirs(thumbnails_dir, exist_ok=True)
|
|
|
|
# Generate thumbnail filename
|
|
pdf_filename = os.path.basename(pdf_path)
|
|
thumb_filename = os.path.splitext(pdf_filename)[0] + '_thumb.png'
|
|
thumb_path = os.path.join(thumbnails_dir, thumb_filename)
|
|
|
|
# Resize and save
|
|
img = images[0]
|
|
img.thumbnail(size, Image.Resampling.LANCZOS)
|
|
img.save(thumb_path, 'PNG', optimize=True)
|
|
|
|
# Return relative path
|
|
rel_path = os.path.join('thumbnails', secure_filename(set_number), thumb_filename)
|
|
return rel_path
|
|
|
|
except ImportError:
|
|
current_app.logger.warning('Neither PyMuPDF nor pdf2image available for PDF thumbnails')
|
|
return None
|
|
|
|
except Exception as e:
|
|
current_app.logger.error(f'PDF thumbnail creation failed: {str(e)}')
|
|
return None
|
|
|
|
@staticmethod
|
|
def delete_file(file_path: str) -> bool:
|
|
"""
|
|
Delete a file from storage.
|
|
|
|
Args:
|
|
file_path: Relative path to the file
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
full_path = os.path.join(current_app.config['UPLOAD_FOLDER'], file_path)
|
|
|
|
if os.path.exists(full_path):
|
|
os.remove(full_path)
|
|
|
|
# Also delete thumbnail if it exists
|
|
base, ext = os.path.splitext(full_path)
|
|
thumb_path = f"{base}_thumb{ext}"
|
|
if os.path.exists(thumb_path):
|
|
os.remove(thumb_path)
|
|
|
|
return True
|
|
return False
|
|
|
|
except Exception as e:
|
|
current_app.logger.error(f'File deletion failed: {str(e)}')
|
|
return False
|
|
|
|
@staticmethod
|
|
def get_directory_size(set_number: str) -> Tuple[int, int]:
|
|
"""
|
|
Get total size of all files for a specific set.
|
|
|
|
Args:
|
|
set_number: LEGO set number
|
|
|
|
Returns:
|
|
Tuple of (total_size_bytes, file_count)
|
|
"""
|
|
total_size = 0
|
|
file_count = 0
|
|
|
|
for subdir in ['pdfs', 'images']:
|
|
set_dir = os.path.join(
|
|
current_app.config['UPLOAD_FOLDER'],
|
|
subdir,
|
|
secure_filename(set_number)
|
|
)
|
|
|
|
if os.path.exists(set_dir):
|
|
for filename in os.listdir(set_dir):
|
|
if not filename.endswith('_thumb'):
|
|
filepath = os.path.join(set_dir, filename)
|
|
if os.path.isfile(filepath):
|
|
total_size += os.path.getsize(filepath)
|
|
file_count += 1
|
|
|
|
return total_size, file_count
|