import hashlib from shutil import COPY_BUFSIZE # type: ignore import blurhash # type: ignore from fastapi import UploadFile from loguru import logger from PIL import Image from sqlalchemy import select from app import activitypub as ap from app import models from app.config import BASE_URL from app.config import ROOT_DIR from app.database import AsyncSession UPLOAD_DIR = ROOT_DIR / "data" / "uploads" async def save_upload(db_session: AsyncSession, f: UploadFile) -> models.Upload: # Compute the hash h = hashlib.blake2b(digest_size=32) while True: buf = f.file.read(COPY_BUFSIZE) if not buf: break h.update(buf) content_hash = h.hexdigest() f.file.seek(0) existing_upload = ( await db_session.execute( select(models.Upload).where(models.Upload.content_hash == content_hash) ) ).scalar_one_or_none() if existing_upload: logger.info(f"Upload with {content_hash=} already exists") return existing_upload logger.info(f"Creating new Upload with {content_hash=}") dest_filename = UPLOAD_DIR / content_hash has_thumbnail = False image_blurhash = None width = None height = None if f.content_type.startswith("image"): image_blurhash = blurhash.encode(f.file, x_components=4, y_components=3) f.file.seek(0) with Image.open(f.file) as original_image: destination_image = Image.new( original_image.mode, original_image.size, ) destination_image.putdata(original_image.getdata()) destination_image.save( dest_filename, format=original_image.format, ) try: width, height = original_image.size original_image.thumbnail((740, 740)) original_image.save( UPLOAD_DIR / f"{content_hash}_resized", format=original_image.format, ) except Exception: logger.exception( f"Failed to created thumbnail for {f.filename}/{content_hash}" ) else: has_thumbnail = True logger.info("Thumbnail generated") else: with open(dest_filename, "wb") as dest: while True: buf = f.file.read(COPY_BUFSIZE) if not buf: break dest.write(buf) new_upload = models.Upload( content_type=f.content_type, content_hash=content_hash, has_thumbnail=has_thumbnail, blurhash=image_blurhash, width=width, height=height, ) db_session.add(new_upload) await db_session.commit() return new_upload def upload_to_attachment(upload: models.Upload, filename: str) -> ap.RawObject: extra_attachment_fields = {} if upload.blurhash: extra_attachment_fields.update( { "blurhash": upload.blurhash, "height": upload.height, "width": upload.width, } ) return { "type": "Document", "mediaType": upload.content_type, "name": filename, "url": BASE_URL + f"/attachments/{upload.content_hash}/{filename}", **extra_attachment_fields, }