115 lines
3.2 KiB
Python
115 lines
3.2 KiB
Python
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 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 Session
|
|
|
|
UPLOAD_DIR = ROOT_DIR / "data" / "uploads"
|
|
|
|
|
|
def save_upload(db: Session, 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 = (
|
|
db.query(models.Upload)
|
|
.filter(models.Upload.content_hash == content_hash)
|
|
.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.add(new_upload)
|
|
db.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,
|
|
}
|