109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
|
import re
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
IMAGE_STATUS_PENDING = "pending"
|
||
|
|
IMAGE_STATUS_PROCESSING = "processing"
|
||
|
|
IMAGE_STATUS_COMPLETED = "completed"
|
||
|
|
IMAGE_STATUS_FAILED = "failed"
|
||
|
|
|
||
|
|
VALID_IMAGE_STATUSES = {
|
||
|
|
IMAGE_STATUS_PENDING,
|
||
|
|
IMAGE_STATUS_PROCESSING,
|
||
|
|
IMAGE_STATUS_COMPLETED,
|
||
|
|
IMAGE_STATUS_FAILED,
|
||
|
|
}
|
||
|
|
|
||
|
|
_PLACEHOLDER_DESCRIPTION_RE = re.compile(r"\{\{\{\{IMAGE:(.*?)\}\}\}\}|\{\{IMAGE:(.*?)\}\}")
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_image_asset(asset: dict[str, Any] | None) -> dict[str, Any] | None:
|
||
|
|
if not isinstance(asset, dict):
|
||
|
|
return None
|
||
|
|
|
||
|
|
placeholder = _as_non_empty_string(asset.get("placeholder"))
|
||
|
|
description = _as_non_empty_string(asset.get("description")) or _extract_description_from_placeholder(
|
||
|
|
placeholder
|
||
|
|
)
|
||
|
|
if not placeholder or not description:
|
||
|
|
return None
|
||
|
|
|
||
|
|
normalized = dict(asset)
|
||
|
|
normalized["index"] = _coerce_int(asset.get("index"), default=0)
|
||
|
|
normalized["placeholder"] = placeholder
|
||
|
|
normalized["description"] = description
|
||
|
|
|
||
|
|
status = _as_non_empty_string(asset.get("status")) or IMAGE_STATUS_PENDING
|
||
|
|
if status not in VALID_IMAGE_STATUSES:
|
||
|
|
normalized["status"] = IMAGE_STATUS_FAILED
|
||
|
|
normalized["error"] = asset.get("error") or f"invalid image status: {status}"
|
||
|
|
return normalized
|
||
|
|
|
||
|
|
normalized["status"] = status
|
||
|
|
normalized["prompt"] = _as_optional_string(asset.get("prompt"))
|
||
|
|
normalized["url"] = _as_optional_string(asset.get("url"))
|
||
|
|
normalized["storage_key"] = _as_optional_string(asset.get("storage_key"))
|
||
|
|
normalized["provider"] = _as_optional_string(asset.get("provider"))
|
||
|
|
normalized["style"] = _as_optional_string(asset.get("style"))
|
||
|
|
normalized["size"] = _as_optional_string(asset.get("size"))
|
||
|
|
normalized["error"] = _as_optional_string(asset.get("error"))
|
||
|
|
normalized["created_at"] = _as_optional_string(asset.get("created_at"))
|
||
|
|
normalized["updated_at"] = _as_optional_string(asset.get("updated_at"))
|
||
|
|
|
||
|
|
if normalized["status"] == IMAGE_STATUS_COMPLETED and not (
|
||
|
|
normalized["url"] or normalized["storage_key"]
|
||
|
|
):
|
||
|
|
normalized["status"] = IMAGE_STATUS_FAILED
|
||
|
|
normalized["error"] = normalized["error"] or "missing image url"
|
||
|
|
|
||
|
|
return normalized
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_image_assets(images: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
|
||
|
|
normalized_assets: list[dict[str, Any]] = []
|
||
|
|
for item in images or []:
|
||
|
|
normalized = normalize_image_asset(item)
|
||
|
|
if normalized:
|
||
|
|
normalized_assets.append(normalized)
|
||
|
|
return normalized_assets
|
||
|
|
|
||
|
|
|
||
|
|
def completed_image_assets(images: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
|
||
|
|
return [
|
||
|
|
asset
|
||
|
|
for asset in normalize_image_assets(images)
|
||
|
|
if asset.get("status") == IMAGE_STATUS_COMPLETED
|
||
|
|
and (asset.get("storage_key") or asset.get("url"))
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def _as_non_empty_string(value: Any) -> str | None:
|
||
|
|
if isinstance(value, str):
|
||
|
|
stripped = value.strip()
|
||
|
|
return stripped or None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _as_optional_string(value: Any) -> str | None:
|
||
|
|
if value is None:
|
||
|
|
return None
|
||
|
|
if isinstance(value, str):
|
||
|
|
return value
|
||
|
|
return str(value)
|
||
|
|
|
||
|
|
|
||
|
|
def _coerce_int(value: Any, default: int) -> int:
|
||
|
|
try:
|
||
|
|
return int(value)
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
return default
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_description_from_placeholder(placeholder: str | None) -> str | None:
|
||
|
|
if not placeholder:
|
||
|
|
return None
|
||
|
|
match = _PLACEHOLDER_DESCRIPTION_RE.fullmatch(placeholder.strip())
|
||
|
|
if not match:
|
||
|
|
return None
|
||
|
|
description = (match.group(1) or match.group(2) or "").strip()
|
||
|
|
return description or None
|