lightning / helper /assets.py
R.C.M.
Improved structure
99d08e9
Raw
History Blame
2.29 kB
# app.py
import base64
import os
import uuid
from pathlib import Path
from typing import Dict
from PIL import Image
from io import BytesIO
from fastapi import FastAPI, HTTPException, APIRouter
from fastapi.responses import FileResponse
BASE_DIR = Path(__file__).parent
ASSETS_DIR = BASE_DIR / "cdn-assets"
ASSETS_DIR.mkdir(exist_ok=True)
image_registry: Dict[str, Path] = {}
asset_router = APIRouter(prefix="/asset-cdn", tags=["asset-cdn"])
def save_base64_image(b64_data: str) -> str:
"""
Decodes a base64 string, converts the image to PNG, stores it in /cdn-assets and returns the
unique image_id that can be used in the public URL.
"""
if "," in b64_data:
_, b64_data = b64_data.split(",", 1)
try:
decoded = base64.b64decode(b64_data)
except (base64.binascii.Error, ValueError) as exc:
raise ValueError("Invalid base64 data") from exc
image = Image.open(BytesIO(decoded))
image_id = uuid.uuid4().hex
file_path = ASSETS_DIR / f"{image_id}.png"
image.save(file_path, "PNG")
image_registry[image_id] = file_path
return image_id
@asset_router.get("/assets/{image_id}")
def get_image(image_id: str):
"""
FastAPI route that streams the requested image to the client.
If the image was cleaned up or never existed, a 404 is returned.
"""
file_path = image_registry.get(image_id)
if not file_path or not file_path.exists():
raise HTTPException(status_code=404, detail="Image not found")
return FileResponse(path=file_path, media_type="image/png")
def cleanup_image(image_id: str) -> bool:
"""
Deletes the image file and removes it from the registry.
Returns True if the image existed and was removed, otherwise False.
"""
file_path = image_registry.pop(image_id, None)
if file_path and file_path.exists():
try:
file_path.unlink()
except Exception:
image_registry[image_id] = file_path
return False
return True
return False
def is_base64_image(value: str) -> bool:
if not isinstance(value, str):
return False
if value.startswith("http://") or value.startswith("https://"):
return False
if value.startswith("data:image"):
return True
return False