Spaces:
Running
Running
| # utils/image_handler.py | |
| from __future__ import annotations | |
| from pathlib import Path | |
| from typing import Optional, Tuple, Union, IO, List, Any | |
| from uuid import uuid4 | |
| import inspect | |
| import os | |
| # IMPORTANT: import FastAPI's UploadFile *just* for type hints; we won't rely on isinstance checks | |
| try: | |
| from fastapi import UploadFile as FastAPIUploadFile # type: ignore | |
| except Exception: # pragma: no cover | |
| FastAPIUploadFile = Any # fallback for typing | |
| UPLOAD_DIR = Path("uploads") | |
| UPLOAD_DIR.mkdir(parents=True, exist_ok=True) | |
| def _find_by_uuid(file_uuid: str) -> Optional[Path]: | |
| matches = list(UPLOAD_DIR.glob(f"{file_uuid}.*")) | |
| return matches[0] if matches else None | |
| async def _read_any(upload: Any) -> bytes: | |
| """ | |
| Read bytes from: | |
| - Starlette/FastAPI UploadFile (async .read) | |
| - File-like object with sync .read | |
| - Bytes/bytearray | |
| """ | |
| # Starlette/FastAPI UploadFile or similar (has filename and an async read) | |
| if hasattr(upload, "filename") and hasattr(upload, "read"): | |
| if inspect.iscoroutinefunction(upload.read): | |
| data = await upload.read() | |
| else: | |
| data = upload.read() | |
| if isinstance(data, memoryview): | |
| data = data.tobytes() | |
| if not isinstance(data, (bytes, bytearray)): | |
| raise TypeError(f"Expected bytes from upload.read(), got {type(data)}") | |
| return bytes(data) | |
| # File-like with sync read() | |
| if hasattr(upload, "read") and callable(upload.read): | |
| data = upload.read() | |
| if inspect.isawaitable(data): # some odd wrappers might return awaitable | |
| data = await data | |
| if isinstance(data, memoryview): | |
| data = data.tobytes() | |
| if not isinstance(data, (bytes, bytearray)): | |
| raise TypeError(f"Expected bytes from file-like read(), got {type(data)}") | |
| return bytes(data) | |
| # Raw bytes | |
| if isinstance(upload, (bytes, bytearray, memoryview)): | |
| return bytes(upload) | |
| raise TypeError(f"Unsupported upload type: {type(upload)}") | |
| def _infer_suffix(upload: Any, default: str = ".png") -> str: | |
| name = None | |
| if hasattr(upload, "filename"): | |
| name = getattr(upload, "filename") or None | |
| return (Path(name).suffix.lower() if name else "") or default | |
| async def save_uploaded_file( | |
| upload: Union[Any, bytes, IO[bytes], tuple, list], | |
| file_uuid: Optional[str] = None, | |
| ) -> Tuple[Path, str, bool]: | |
| """ | |
| Save under uploads/<UUID>.<ext>. If same UUID exists, do not rewrite. | |
| Returns: (path, uuid, deduped) | |
| """ | |
| # Unwrap common mistakes like (UploadFile,) or [UploadFile] | |
| if isinstance(upload, (tuple, list)) and upload: | |
| upload = upload[0] | |
| file_uuid = (file_uuid or str(uuid4())).strip() | |
| # Idempotent by UUID: if already saved, don't write again | |
| existing = _find_by_uuid(file_uuid) | |
| if existing and existing.exists(): | |
| return existing, file_uuid, True | |
| suffix = _infer_suffix(upload, default=".png") | |
| dest = UPLOAD_DIR / f"{file_uuid}{suffix}" | |
| data = await _read_any(upload) | |
| with open(dest, "wb") as out: | |
| out.write(data) | |
| return dest, file_uuid, False | |
| def load_image(path: Path): | |
| from PIL import Image | |
| return Image.open(path) | |
| def is_gif(path: Path) -> bool: | |
| return path.suffix.lower() == ".gif" | |