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)