microblog.pub/app/uploads.py

120 lines
3.4 KiB
Python
Raw Normal View History

2022-06-23 19:07:20 +00: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 06:56:39 +00:00
from sqlalchemy import select
2022-06-23 19:07:20 +00: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 18:43:17 +00:00
from app.database import AsyncSession
2022-06-23 19:07:20 +00:00
UPLOAD_DIR = ROOT_DIR / "data" / "uploads"
2022-06-29 18:43:17 +00:00
async def save_upload(db_session: AsyncSession, f: UploadFile) -> models.Upload:
2022-06-23 19:07:20 +00: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 20:07:18 +00:00
f.file.seek(0)
2022-06-23 19:07:20 +00:00
2022-06-29 18:43:17 +00:00
existing_upload = (
await db_session.execute(
select(models.Upload).where(models.Upload.content_hash == content_hash)
)
2022-06-29 06:56:39 +00:00
).scalar_one_or_none()
2022-06-23 19:07:20 +00: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 20:07:18 +00: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 19:07:20 +00:00
)
2022-06-23 20:07:18 +00: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 19:07:20 +00: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 18:43:17 +00:00
db_session.add(new_upload)
await db_session.commit()
2022-06-23 19:07:20 +00:00
return new_upload
2022-07-21 20:43:06 +00:00
def upload_to_attachment(
upload: models.Upload,
filename: str,
alt_text: str | None,
) -> ap.RawObject:
2022-06-23 19:07:20 +00:00
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,
2022-07-21 20:43:06 +00:00
"name": alt_text or filename,
2022-06-23 20:07:18 +00:00
"url": BASE_URL + f"/attachments/{upload.content_hash}/{filename}",
2022-06-23 19:07:20 +00:00
**extra_attachment_fields,
}