microblog.pub/app/uploads.py

116 lines
3.3 KiB
Python
Raw Normal View History

2022-06-23 14:07:20 -05:00
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
2022-06-29 01:56:39 -05:00
from sqlalchemy import select
2022-06-23 14:07:20 -05:00
from app import activitypub as ap
from app import models
from app.config import BASE_URL
from app.config import ROOT_DIR
2022-06-29 13:43:17 -05:00
from app.database import AsyncSession
2022-06-23 14:07:20 -05:00
UPLOAD_DIR = ROOT_DIR / "data" / "uploads"
2022-06-29 13:43:17 -05:00
async def save_upload(db_session: AsyncSession, f: UploadFile) -> models.Upload:
2022-06-23 14:07:20 -05:00
# 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()
2022-06-23 15:07:18 -05:00
f.file.seek(0)
2022-06-23 14:07:20 -05:00
2022-06-29 13:43:17 -05:00
existing_upload = (
await db_session.execute(
select(models.Upload).where(models.Upload.content_hash == content_hash)
)
2022-06-29 01:56:39 -05:00
).scalar_one_or_none()
2022-06-23 14:07:20 -05:00
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"):
2022-06-23 15:07:18 -05:00
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,
2022-06-23 14:07:20 -05:00
)
2022-06-23 15:07:18 -05:00
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)
2022-06-23 14:07:20 -05:00
new_upload = models.Upload(
content_type=f.content_type,
content_hash=content_hash,
has_thumbnail=has_thumbnail,
blurhash=image_blurhash,
width=width,
height=height,
)
2022-06-29 13:43:17 -05:00
db_session.add(new_upload)
await db_session.commit()
2022-06-23 14:07:20 -05:00
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,
2022-06-23 15:07:18 -05:00
"url": BASE_URL + f"/attachments/{upload.content_hash}/{filename}",
2022-06-23 14:07:20 -05:00
**extra_attachment_fields,
}