Spaces:
Paused
Paused
File size: 3,098 Bytes
8d1819a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import os
import io
import base64
from PIL import Image
from typing import Dict, List, Optional, Tuple
from werkzeug.utils import secure_filename
from python.helpers.print_style import PrintStyle
class AttachmentManager:
ALLOWED_EXTENSIONS = {
'image': {'jpg', 'jpeg', 'png', 'bmp'},
'code': {'py', 'js', 'sh', 'html', 'css'},
'document': {'md', 'pdf', 'txt', 'csv', 'json'}
}
def __init__(self, work_dir: str):
self.work_dir = work_dir
os.makedirs(work_dir, exist_ok=True)
def is_allowed_file(self, filename: str) -> bool:
ext = self.get_file_extension(filename)
all_allowed = set().union(*self.ALLOWED_EXTENSIONS.values())
return ext in all_allowed
def get_file_type(self, filename: str) -> str:
ext = self.get_file_extension(filename)
for file_type, extensions in self.ALLOWED_EXTENSIONS.items():
if ext in extensions:
return file_type
return 'unknown'
@staticmethod
def get_file_extension(filename: str) -> str:
return filename.rsplit('.', 1)[1].lower() if '.' in filename else ''
def validate_mime_type(self, file) -> bool:
try:
mime_type = file.content_type
return mime_type.split('/')[0] in ['image', 'text', 'application']
except AttributeError:
return False
def save_file(self, file, filename: str) -> Tuple[str, Dict]:
"""Save file and return path and metadata"""
try:
filename = secure_filename(filename)
if not filename:
raise ValueError("Invalid filename")
file_path = os.path.join(self.work_dir, filename)
file_type = self.get_file_type(filename)
metadata = {
'filename': filename,
'type': file_type,
'extension': self.get_file_extension(filename),
'preview': None
}
# Save file
file.save(file_path)
# Generate preview for images
if file_type == 'image':
metadata['preview'] = self.generate_image_preview(file_path)
return file_path, metadata
except Exception as e:
PrintStyle.error(f"Error saving file {filename}: {e}")
return None, {} # type: ignore
def generate_image_preview(self, image_path: str, max_size: int = 800) -> Optional[str]:
try:
with Image.open(image_path) as img:
# Convert image if needed
if img.mode in ('RGBA', 'P'):
img = img.convert('RGB')
# Resize for preview
img.thumbnail((max_size, max_size))
# Save to buffer
buffer = io.BytesIO()
img.save(buffer, format="JPEG", quality=70, optimize=True)
# Convert to base64
return base64.b64encode(buffer.getvalue()).decode('utf-8')
except Exception as e:
PrintStyle.error(f"Error generating preview for {image_path}: {e}")
return None
|