File size: 3,007 Bytes
609d54c 18d9ea5 |
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 94 95 96 97 98 99 100 101 |
import shutil
from custom_logger import logger_config
import os
from pathlib import Path
class MainBase:
def __init__(self):
self.max_file_size = 10 * 1024 * 1024
def upload_validate(self, image):
# Reset file pointer to beginning before reading
image.file.seek(0)
contents = image.file.read()
# Reset file pointer again for later use
image.file.seek(0)
if len(contents) > self.max_file_size:
raise ValueError(f"{image.filename} is too large. Max size: {self.max_file_size // (1024*1024)}MB.")
def upload(self, id, image):
self.upload_validate(image)
# Create full file path with filename
file_path = os.path.join(self.input_dir, id)
with open(file_path, "wb") as buffer:
shutil.copyfileobj(image.file, buffer)
logger_config.success(f"{image.filename} uploaded successfully to {file_path}")
def _generate_output_path(self, output_format=None) -> str:
if output_format:
new_ext = self.supported_formats.get(output_format, output_format).lower()
name_without_ext = os.path.splitext(self.input_file_name)[0]
return f"{self.output_dir}/{name_without_ext}.{new_ext}"
return f"{self.output_dir}/{self.input_file_name}"
def download_url(self, file_id):
self.delete([file_id], only_input=True)
file_path = f'{self.output_dir}/{file_id}'
if os.path.exists(file_path):
return file_path
raise ValueError("File not found")
def delete(self, ids, only_input=False):
for file_id in ids:
name_without_ext = os.path.splitext(file_id)[0]
for filename in os.listdir(self.input_dir):
if name_without_ext in filename:
file_path = os.path.join(self.input_dir, filename)
path_obj = Path(file_path)
path_obj.unlink(missing_ok=True)
if not only_input:
for filename in os.listdir(self.output_dir):
if name_without_ext in filename:
file_path = os.path.join(self.output_dir, filename)
path_obj = Path(file_path)
path_obj.unlink(missing_ok=True)
def cleanup_old_files(self, max_age_days=2):
"""
Remove files older than max_age_days from input and output directories.
Args:
max_age_days: Maximum age in days before files are deleted (default: 2)
Returns:
Number of files deleted
"""
import time
deleted_count = 0
max_age_seconds = max_age_days * 24 * 60 * 60
current_time = time.time()
directories = [self.input_dir, self.output_dir]
for directory in directories:
if not os.path.exists(directory):
continue
for filename in os.listdir(directory):
file_path = os.path.join(directory, filename)
# Skip directories
if os.path.isdir(file_path):
continue
try:
file_mtime = os.path.getmtime(file_path)
file_age = current_time - file_mtime
if file_age > max_age_seconds:
os.unlink(file_path)
deleted_count += 1
logger_config.info(f"Cleaned up old file: {filename}")
except Exception as e:
logger_config.warning(f"Failed to cleanup {filename}: {e}")
return deleted_count |