File size: 1,948 Bytes
609d54c |
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 |
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) |