211 lines
7.7 KiB
Python
211 lines
7.7 KiB
Python
"""Safe local export helpers for AI image studio final selections."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import shutil
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
|
|
from . import appconfig, image_studio
|
|
|
|
|
|
EXISTING_FAIL = "fail"
|
|
EXISTING_OVERWRITE_MANAGED = "overwrite_managed"
|
|
EXISTING_TIMESTAMP = "timestamp"
|
|
EXISTING_MODES = {EXISTING_FAIL, EXISTING_OVERWRITE_MANAGED, EXISTING_TIMESTAMP}
|
|
|
|
|
|
class ImageStudioExportError(RuntimeError):
|
|
"""Raised when AI studio selections cannot be exported safely."""
|
|
|
|
|
|
class ExportTargetExistsError(ImageStudioExportError):
|
|
"""Raised when the default target directory already exists."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExportedFile:
|
|
selection_type: str
|
|
asset_id: int
|
|
source_path: str
|
|
output_path: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExportResult:
|
|
target_dir: str
|
|
files: tuple[ExportedFile, ...]
|
|
main_count: int
|
|
detail_count: int
|
|
existing_mode: str
|
|
|
|
|
|
def target_dir_for_project(project, parent_dir, suffix=None):
|
|
parent = _existing_parent_dir(parent_dir)
|
|
item = _safe_item_id(getattr(project, "item_id", ""))
|
|
dirname = item if not suffix else f"{item}_{suffix}"
|
|
return os.path.abspath(os.path.join(parent, dirname))
|
|
|
|
|
|
def export_project_selection(
|
|
project_id,
|
|
parent_dir,
|
|
*,
|
|
existing_mode=EXISTING_FAIL,
|
|
path=None,
|
|
config=None,
|
|
timestamp=None,
|
|
):
|
|
"""Export current main/detail selections as ordered JPEG files."""
|
|
|
|
mode = str(existing_mode or EXISTING_FAIL)
|
|
if mode not in EXISTING_MODES:
|
|
raise ImageStudioExportError("导出目录处理方式无效")
|
|
cfg = appconfig.load_config() if config is None else config
|
|
database_path = path or appconfig.db_path(cfg)
|
|
project = image_studio.get_project(project_id, path=database_path)
|
|
if project is None:
|
|
raise ImageStudioExportError("AI工场项目不存在")
|
|
parent = _existing_parent_dir(parent_dir)
|
|
planned = _planned_files(project, database_path)
|
|
if not planned:
|
|
raise ImageStudioExportError("主图和详情图终选都为空，不能导出")
|
|
target = _choose_target_dir(project, parent, mode, timestamp=timestamp)
|
|
quality = int(appconfig.ai_config(cfg).get("jpg_quality", 90) or 90)
|
|
staging = os.path.join(parent, f".{_safe_item_id(project.item_id)}_staging_{uuid.uuid4().hex}")
|
|
try:
|
|
os.makedirs(staging, exist_ok=False)
|
|
staged = []
|
|
for item in planned:
|
|
output_path = os.path.join(staging, item["filename"])
|
|
_save_jpeg(item["source_path"], output_path, quality)
|
|
staged.append(ExportedFile(item["selection_type"], item["asset_id"], item["source_path"], output_path))
|
|
os.makedirs(target, exist_ok=True)
|
|
if os.path.exists(target) and mode == EXISTING_OVERWRITE_MANAGED:
|
|
_remove_managed_exports(target, project.item_id)
|
|
final_files = []
|
|
for staged_file in staged:
|
|
final_path = os.path.join(target, os.path.basename(staged_file.output_path))
|
|
os.replace(staged_file.output_path, final_path)
|
|
final_files.append(
|
|
ExportedFile(
|
|
staged_file.selection_type,
|
|
staged_file.asset_id,
|
|
staged_file.source_path,
|
|
final_path,
|
|
)
|
|
)
|
|
return ExportResult(
|
|
target_dir=target,
|
|
files=tuple(final_files),
|
|
main_count=sum(1 for item in final_files if item.selection_type == "main"),
|
|
detail_count=sum(1 for item in final_files if item.selection_type == "detail"),
|
|
existing_mode=mode,
|
|
)
|
|
except Exception as exc:
|
|
if isinstance(exc, ImageStudioExportError):
|
|
raise
|
|
raise ImageStudioExportError(f"导出终选图片失败：{exc}") from exc
|
|
finally:
|
|
if os.path.isdir(staging):
|
|
shutil.rmtree(staging, ignore_errors=True)
|
|
|
|
|
|
def _choose_target_dir(project, parent_dir, existing_mode, timestamp=None):
|
|
target = target_dir_for_project(project, parent_dir)
|
|
if not os.path.exists(target):
|
|
return target
|
|
if not os.path.isdir(target):
|
|
raise ImageStudioExportError(f"导出目标已存在但不是目录：{target}")
|
|
if existing_mode == EXISTING_FAIL:
|
|
raise ExportTargetExistsError(f"商品目录已存在：{target}")
|
|
if existing_mode == EXISTING_OVERWRITE_MANAGED:
|
|
return target
|
|
stamp = timestamp or datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
base = target_dir_for_project(project, parent_dir, suffix=stamp)
|
|
candidate = base
|
|
index = 2
|
|
while os.path.exists(candidate):
|
|
candidate = f"{base}_{index}"
|
|
index += 1
|
|
return candidate
|
|
|
|
|
|
def _planned_files(project, db_path):
|
|
item = _safe_item_id(project.item_id)
|
|
planned = []
|
|
for selection_type, label in (("main", "主图"), ("detail", "详情图")):
|
|
selections = image_studio.list_selections(project.id, selection_type, path=db_path)
|
|
for index, selection in enumerate(selections, start=1):
|
|
asset = image_studio.get_asset(selection.asset_id, path=db_path)
|
|
if asset is None:
|
|
raise ImageStudioExportError(f"终选照片不存在：#{selection.asset_id}")
|
|
source_path = str(asset.local_path or "")
|
|
if not source_path or not os.path.isfile(source_path):
|
|
raise ImageStudioExportError(f"终选照片本地文件缺失：#{asset.id}")
|
|
planned.append(
|
|
{
|
|
"selection_type": selection_type,
|
|
"asset_id": int(asset.id),
|
|
"source_path": os.path.abspath(source_path),
|
|
"filename": f"{item}_{label}_{index}.jpg",
|
|
}
|
|
)
|
|
return planned
|
|
|
|
|
|
def _save_jpeg(source_path, output_path, quality):
|
|
try:
|
|
from PIL import Image
|
|
|
|
with Image.open(source_path) as image:
|
|
if image.mode in {"RGBA", "LA"} or (
|
|
image.mode == "P" and "transparency" in getattr(image, "info", {})
|
|
):
|
|
rgba = image.convert("RGBA")
|
|
background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
|
|
background.alpha_composite(rgba)
|
|
final_image = background.convert("RGB")
|
|
else:
|
|
final_image = image.convert("RGB")
|
|
final_image.save(output_path, format="JPEG", quality=max(1, min(95, int(quality or 90))))
|
|
with Image.open(output_path) as check:
|
|
check.verify()
|
|
except Exception as exc:
|
|
raise ImageStudioExportError(f"JPEG 转码失败：{exc}") from exc
|
|
|
|
|
|
def _remove_managed_exports(target_dir, item_id):
|
|
item = re.escape(_safe_item_id(item_id))
|
|
pattern = re.compile(rf"^{item}_(主图|详情图)_\d+\.jpg$", re.IGNORECASE)
|
|
for filename in os.listdir(target_dir):
|
|
if not pattern.match(filename):
|
|
continue
|
|
path = os.path.join(target_dir, filename)
|
|
if os.path.isfile(path):
|
|
os.remove(path)
|
|
|
|
|
|
def _existing_parent_dir(parent_dir):
|
|
parent = os.path.abspath(str(parent_dir or ""))
|
|
if not os.path.isdir(parent):
|
|
raise ImageStudioExportError(f"导出父目录不存在：{parent}")
|
|
return parent
|
|
|
|
|
|
def _safe_item_id(item_id):
|
|
value = str(item_id or "").strip()
|
|
if not value:
|
|
raise ImageStudioExportError("商品ID不能为空")
|
|
if any(char in value for char in ('/', '\\', os.sep, os.altsep or "\0")):
|
|
raise ImageStudioExportError("商品ID包含非法路径字符")
|
|
if value in {".", ".."}:
|
|
raise ImageStudioExportError("商品ID不能是路径符号")
|
|
safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in value).strip("_")
|
|
if not safe:
|
|
raise ImageStudioExportError("商品ID不能作为目录名")
|
|
return safe
|