From cddadf636495fc240c82d89b46e4d321814851cd Mon Sep 17 00:00:00 2001 From: chengma Date: Sat, 11 Jul 2026 14:26:01 +0800 Subject: [PATCH] feat(ai-studio): export final selections --- app/gui/__init__.py | 1 + app/gui/tabs/image_studio.py | 111 +++++++++++++++- app/gui/workers.py | 40 +++++- app/image_studio_export.py | 210 ++++++++++++++++++++++++++++++ docs/04-architecture.md | 5 +- docs/routes.md | 7 +- docs/tasks/T-593.md | 9 +- tests/test_gui.py | 2 + tests/test_image_studio_export.py | 162 +++++++++++++++++++++++ 9 files changed, 539 insertions(+), 8 deletions(-) create mode 100644 app/image_studio_export.py create mode 100644 tests/test_image_studio_export.py diff --git a/app/gui/__init__.py b/app/gui/__init__.py index 13a0aaa..3821b40 100644 --- a/app/gui/__init__.py +++ b/app/gui/__init__.py @@ -25,6 +25,7 @@ if QT_IMPORT_ERROR is None: CollectWorker, GenerateWorker, ImageStudioDownloadOriginalWorker, + ImageStudioExportWorker, ImageStudioGenerateJobsWorker, ImageStudioPullImagesWorker, WriteBackWorker, diff --git a/app/gui/tabs/image_studio.py b/app/gui/tabs/image_studio.py index f8146a0..86022ec 100644 --- a/app/gui/tabs/image_studio.py +++ b/app/gui/tabs/image_studio.py @@ -7,12 +7,13 @@ import os from PySide6.QtCore import QMimeData from PySide6.QtWidgets import QListWidget, QListWidgetItem -from ... import accounts, appconfig, db, image_studio, prompts +from ... import accounts, appconfig, db, image_studio, image_studio_export, prompts from .. import file_manager from ..widgets import * from ..workers import ( ImageStudioDownloadOriginalWorker as _RealImageStudioDownloadOriginalWorker, ) +from ..workers import ImageStudioExportWorker as _RealImageStudioExportWorker from ..workers import ImageStudioGenerateJobsWorker as _RealImageStudioGenerateJobsWorker from ..workers import ImageStudioPullImagesWorker as _RealImageStudioPullImagesWorker @@ -47,6 +48,15 @@ def ImageStudioGenerateJobsWorker(*args, **kwargs): ) +def ImageStudioExportWorker(*args, **kwargs): + return _call_package_attr( + "ImageStudioExportWorker", + _RealImageStudioExportWorker, + *args, + **kwargs, + ) + + def _drop_event_position(event): if hasattr(event, "position"): return event.position().toPoint() @@ -469,8 +479,20 @@ class ImageStudioTab(QWidget): self.detail_selection_list.setMinimumHeight(96) detail_layout.addWidget(self.detail_selection_label) detail_layout.addWidget(self.detail_selection_list) + action_panel = QWidget() + action_layout = QVBoxLayout(action_panel) + action_layout.setContentsMargins(0, 0, 0, 0) + self.export_button = QPushButton("导出终选") + self.export_button.setObjectName("imageStudioExportButton") + self.export_hint_label = QLabel("可部分导出，不要求主图/详情图满额") + self.export_hint_label.setObjectName("imageStudioExportHintLabel") + self.export_hint_label.setWordWrap(True) + action_layout.addWidget(self.export_button) + action_layout.addWidget(self.export_hint_label) + action_layout.addStretch(1) layout.addWidget(main_panel, 1) layout.addWidget(detail_panel, 1) + layout.addWidget(action_panel, 0) return panel def _connect_signals(self): @@ -491,6 +513,7 @@ class ImageStudioTab(QWidget): self.prompt_edit.textChanged.connect(self._save_project_prompt) self.start_button.clicked.connect(self.start_generation) self.stop_button.clicked.connect(self.stop_generation) + self.export_button.clicked.connect(self.export_selections) def refresh_accounts(self): self.account_combo.clear() @@ -1032,6 +1055,91 @@ class ImageStudioTab(QWidget): self._append_log("[AI工场] 已请求停止，正在等待安全边界") self._status("AI工场生成已请求停止", "warning") + def export_selections(self, checked=False): + if self.current_project is None: + self._message("项目未打开", "请先打开一个AI工场项目。") + return + main_count = len(self._selection_asset_ids("main")) + detail_count = len(self._selection_asset_ids("detail")) + if main_count + detail_count <= 0: + self._message("没有可导出图片", "主图和详情图终选都为空。") + return + parent_dir = QFileDialog.getExistingDirectory(self, "选择导出父目录") + if not parent_dir: + return + mode = self._export_existing_mode(parent_dir, main_count, detail_count) + if mode is None: + return + worker = ImageStudioExportWorker( + self.current_project.id, + parent_dir, + existing_mode=mode, + db_path=self.db_path, + config=self.config, + ) + worker.log.connect(self._append_log) + worker.finished.connect(self._on_export_finished) + worker.failed.connect(self._on_worker_failed) + self._start_worker(worker, "AI工场导出终选") + + def _export_existing_mode(self, parent_dir, main_count, detail_count): + try: + target = image_studio_export.target_dir_for_project(self.current_project, parent_dir) + except Exception as exc: + self._message("导出目录不可用", str(exc)) + return None + if not os.path.exists(target): + return image_studio_export.EXISTING_FAIL + box = QMessageBox(self) + box.setWindowTitle("商品目录已存在") + box.setText( + f"目标目录已存在：{target}\n" + f"本次将导出主图 {main_count} 张、详情图 {detail_count} 张。请选择处理方式。" + ) + overwrite_button = box.addButton("覆盖本软件导出的图片", QMessageBox.AcceptRole) + timestamp_button = box.addButton("新建带时间目录", QMessageBox.ActionRole) + box.addButton("取消", QMessageBox.RejectRole) + box.setDefaultButton(timestamp_button) + box.exec() + clicked = box.clickedButton() + if clicked is overwrite_button: + return image_studio_export.EXISTING_OVERWRITE_MANAGED + if clicked is timestamp_button: + return image_studio_export.EXISTING_TIMESTAMP + return None + + def _on_export_finished(self, summary): + if summary.get("ok") is False: + if self._running_worker is not None: + self._on_worker_failed(-1, summary.get("error") or "导出终选失败") + return + self._finish_worker() + target_dir = summary.get("target_dir") or "" + main_count = int(summary.get("main_count") or 0) + detail_count = int(summary.get("detail_count") or 0) + file_count = int(summary.get("file_count") or 0) + self._append_log(f"[AI工场] 导出完成：主图{main_count}，详情图{detail_count}，文件{file_count}") + self._status(f"AI工场终选已导出：{file_count} 个文件", "success") + self._show_export_success(target_dir, main_count, detail_count, file_count) + + def _show_export_success(self, target_dir, main_count, detail_count, file_count): + box = QMessageBox(self) + box.setWindowTitle("导出完成") + box.setText( + f"已导出 {file_count} 个 JPEG 文件。\n" + f"主图 {main_count} 张，详情图 {detail_count} 张。\n" + f"目录：{target_dir}" + ) + open_button = box.addButton("打开目录", QMessageBox.AcceptRole) + box.addButton("确定", QMessageBox.RejectRole) + box.setDefaultButton(open_button) + box.exec() + if box.clickedButton() is open_button: + try: + file_manager.open_in_file_manager(target_dir) + except Exception as exc: + self._message("打开目录失败", str(exc)) + def _on_generate_progress(self, payload): total = max(1, int(payload.get("total") or self.progress_bar.maximum() or 1)) done = min(total, int(payload.get("done") or 0)) @@ -1096,6 +1204,7 @@ class ImageStudioTab(QWidget): self.aspect_combo.setEnabled(not running) self.start_button.setEnabled(not running) self.stop_button.setEnabled(running) + self.export_button.setEnabled(not running) def _set_account_combo(self, alias): index = self.account_combo.findData(alias) diff --git a/app/gui/workers.py b/app/gui/workers.py index df57dfa..c4650f2 100644 --- a/app/gui/workers.py +++ b/app/gui/workers.py @@ -6,7 +6,7 @@ import datetime as _dt import re import threading -from .. import image_studio, image_studio_generation, image_studio_images +from .. import image_studio, image_studio_export, image_studio_generation, image_studio_images from .widgets import * @@ -192,6 +192,44 @@ class ImageStudioGenerateJobsWorker(BaseWorker): return summary +class ImageStudioExportWorker(BaseWorker): + """Export AI studio final selections to local JPEG files.""" + + def __init__( + self, + project_id, + parent_dir, + *, + existing_mode=image_studio_export.EXISTING_FAIL, + db_path=None, + config=None, + ): + super().__init__() + self.project_id = int(project_id) + self.parent_dir = parent_dir + self.existing_mode = existing_mode + self.db_path = db_path + self.config = config + + def execute(self): + self.log.emit("[AI工场] 导出终选：开始") + result = image_studio_export.export_project_selection( + self.project_id, + self.parent_dir, + existing_mode=self.existing_mode, + path=self.db_path, + config=self.config, + ) + self.log.emit("[AI工场] 导出终选：成功") + return { + "target_dir": result.target_dir, + "main_count": result.main_count, + "detail_count": result.detail_count, + "file_count": len(result.files), + "existing_mode": result.existing_mode, + } + + def _generation_mode_label(mode): mode = appconfig.normalize_generate_mode(mode) return { diff --git a/app/image_studio_export.py b/app/image_studio_export.py new file mode 100644 index 0000000..9c3a83e --- /dev/null +++ b/app/image_studio_export.py @@ -0,0 +1,210 @@ +"""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 diff --git a/docs/04-architecture.md b/docs/04-architecture.md index 0ee9741..86daa2f 100644 --- a/docs/04-architecture.md +++ b/docs/04-architecture.md @@ -27,7 +27,8 @@ GUI（PySide6 QTabWidget，6 Tab） ├── ai 文本生成（提示词+旧标题→新标题）/ 图像生成（提示词+旧封面→新封面） ├── image_studio AI工场项目/资产/job/终选顺序数据服务 ├── image_studio_images 远程原图安全下载、缩略图、原图落盘 - └── image_studio_generation cmhub 托管多图异步 submit/poll/download 编排 + ├── image_studio_generation cmhub 托管多图异步 submit/poll/download 编排 + └── image_studio_export 终选图片本地 JPEG 转码、目录安全导出 | v Google Chrome（每账号独立 --user-data-dir + --remote-debugging-port） + AI 服务（默认 cmhub 网关；direct 仅内部兼容/回滚） @@ -39,7 +40,7 @@ Shopee 卖家中心页面 / 本地图片目录 真实组件： - GUI 入口：根目录 `main.py` 调用 `app/gui/` 包（PySide6 + `QMainWindow` + `QTabWidget`，6 Tab）；包入口 `app/gui/__init__.py` 提供 `main()` 并兼容 `from app import gui` / `from app.gui import MainWindow`；也支持 `python -m app`。 -- 核心模块统一放在正式代码包 `app/`：`appconfig.py`、`db.py`、`excel.py`、`config.py`、`accounts.py`、`chrome.py`、`editor.py`、`workers.py`、`ai.py`、`prompts.py`、`image_studio.py`、`image_studio_images.py`、`image_studio_generation.py`；CDP 底座迁入 `app/cdp.py`（当前根目录 `cdp.py` 为已验证来源）。 +- 核心模块统一放在正式代码包 `app/`：`appconfig.py`、`db.py`、`excel.py`、`config.py`、`accounts.py`、`chrome.py`、`editor.py`、`workers.py`、`ai.py`、`prompts.py`、`image_studio.py`、`image_studio_images.py`、`image_studio_generation.py`、`image_studio_export.py`；CDP 底座迁入 `app/cdp.py`（当前根目录 `cdp.py` 为已验证来源）。 - 已验证脚本（重构进模块）：`prototypes/demo.py`、`prototypes/set_title.py`、`prototypes/set_cover.py`、`prototypes/get_title.py`、`prototypes/cookies.py`、`prototypes/inspect_images.py`、`prototypes/grab.py`。 - 外部依赖：本机 Google Chrome；Shopee；AI 服务（文本+图像；普通产品默认 cmhub 网关，由 `data/config.json` 的 `ai.cmhub` + `data/config/cmhub.json` 配置；direct 直连模型清单仅作为内部兼容/手工回滚路径保留）；`openpyxl`。 diff --git a/docs/routes.md b/docs/routes.md index bbb2045..0662264 100644 --- a/docs/routes.md +++ b/docs/routes.md @@ -186,7 +186,7 @@ │ [完整提示词输入框] │ │ 类型[主图▼] 数量[4] 比例[1:1▼] cmhub扣点/余额提示 │ │ [开始生成][停止] 进度条 运行日志 │ -│ 底部：主图终选 / 详情图终选（拖入、插入、重排、移出） │ +│ 底部：主图终选 / 详情图终选（拖入、插入、重排、移出） [导出终选] │ └───────────────────────────────────────────────────────────────┘ ``` @@ -200,7 +200,10 @@ - 终选列表内可拖动重排，Delete 或右键「移出终选」只移出终选，不删除照片池资产或本地文件；拖放/移出失败时刷新回 SQLite 中的持久化顺序。 - 主图推荐 1:1；比例不匹配只用黄色轻提示和 tooltip 提醒，不硬拦。文件缺失或尚未下载的照片不能拖入终选。 - 拉主图、下载原图、生图 submit/poll/download 均通过 worker 执行，主线程只刷新 UI；运行中禁用项目切换、模板编辑、源图选择、终选拖放和生成设置，停止为协作式停止。 -- 本小节当前覆盖 T-591/T-592：导出 JPEG 由 T-593 接入，不自动上传或修改蝦皮。 +- 「导出终选」可在主图/详情图未满目标数量时导出当前终选；主图和详情图总数为 0 时阻断。用户选择导出父目录后，程序在其下创建商品 ID 子目录，按终选顺序转码为 `商品ID_主图_1.jpg`、`商品ID_详情图_1.jpg`，透明图铺白底输出真正 JPEG。 +- 商品目录已存在时只提供三选：覆盖本软件导出的图片（仅删除匹配当前商品命名规则的旧主图/详情图，保留用户其它文件）、新建带时间目录、取消；不提供合并，也不递归清空用户目录。 +- 导出前会预检所有终选源文件和图片解码，先写 staging，转码失败不创建商品目录、不留下半套新图；成功后中文提示实际目录和主图/详情图数量，并提供打开目录。 +- 本小节当前覆盖 T-591/T-593：主界面、终选排序和本地导出已接入；不自动上传或修改蝦皮，BYOK/自定义 Provider 仍后置。 ## 流程导航 diff --git a/docs/tasks/T-593.md b/docs/tasks/T-593.md index 79de840..00e6c62 100644 --- a/docs/tasks/T-593.md +++ b/docs/tasks/T-593.md @@ -3,7 +3,7 @@ id: T-593 title: AI工场部分导出、JPEG 转码与已存在商品目录安全处理 phase: 7 deps: [T-592] -status: TODO +status: DONE created: 2026-07-11 --- @@ -39,4 +39,9 @@ created: 2026-07-11 ## 执行记录 -（完成后记录文件安全测试和人工导出验证。） +- 2026-07-11：完成 AI工场终选本地导出。 + - 新增 `app/image_studio_export.py`：读取 AI工场主图/详情图终选顺序，预检源文件和图片解码，按 `商品ID_主图_1.jpg` / `商品ID_详情图_1.jpg` 连续命名转码为真正 JPEG，透明图铺白底。 + - 目录安全：默认目标为用户选择父目录下的 `<商品ID>/`；目录已存在时支持“覆盖本软件导出的图片”（仅删除当前商品受管命名文件，保留用户其它文件）或“新建带时间目录”；不提供合并、不递归清空目录。 + - GUI：⑥ AI工场底部新增「导出终选」按钮，允许部分导出，0 张阻断；导出通过 `ImageStudioExportWorker` 后台执行，成功后中文提示主图/详情图数量、目录，并可打开目录。 + - 更新 `docs/routes.md` 和 `docs/04-architecture.md`，明确导出服务职责和不自动上传蝦皮边界。 + - 验证：主工作区 targeted 测试通过；因无关默认提示词脏文件仍会影响全量 unittest，已在干净 worktree 仅套用 T-593 diff 后运行 `python -m ruff check app tests main.py`、`py -3.10 -m compileall app main.py`、`py -3.10 -m unittest discover -s tests`（367 tests）和 `git diff --check`，全部通过。 diff --git a/tests/test_gui.py b/tests/test_gui.py index 06929f9..7b508cd 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -511,6 +511,8 @@ class GuiTests(TempDirMixin, unittest.TestCase): "导入本地图片", " ".join(button.text() for button in tab.findChildren(gui.QPushButton)), ) + self.assertEqual("导出终选", tab.export_button.text()) + self.assertIn("可部分导出", tab.export_hint_label.text()) self.assertEqual("完整提示词", prompts.load_image_studio_template("工场模板", prompts_dir)) template_index = tab.template_combo.findData("工场模板") self.assertGreaterEqual(template_index, 0) diff --git a/tests/test_image_studio_export.py b/tests/test_image_studio_export.py new file mode 100644 index 0000000..a42e3b3 --- /dev/null +++ b/tests/test_image_studio_export.py @@ -0,0 +1,162 @@ +import io +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(__file__)) + +from _helpers import TempDirMixin + +from app import db, image_studio, image_studio_export + + +class ImageStudioExportTests(TempDirMixin, unittest.TestCase): + def _png_bytes(self, color=(20, 120, 200, 180)): + from PIL import Image + + output = io.BytesIO() + Image.new("RGBA", (24, 24), color).save(output, format="PNG") + return output.getvalue() + + def _write_image(self, path, color=(20, 120, 200, 180)): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "wb") as fh: + fh.write(self._png_bytes(color=color)) + return path + + def _project_with_assets(self, temp_dir): + cfg = { + "db_path": os.path.join(temp_dir, "cmshopee.db"), + "image_dir": os.path.join(temp_dir, "images"), + "ai": {"jpg_quality": 88}, + } + db.init_db(cfg["db_path"]) + project = image_studio.create_or_get_project( + account_alias="alias", + account_slug="alias_slug", + item_id="51100639510", + path=cfg["db_path"], + ) + main = image_studio.add_asset( + project.id, + "generated_main", + local_path=self._write_image(os.path.join(temp_dir, "main.png")), + path=cfg["db_path"], + ) + detail = image_studio.add_asset( + project.id, + "generated_detail", + local_path=self._write_image(os.path.join(temp_dir, "detail.webp"), color=(200, 80, 40, 255)), + path=cfg["db_path"], + ) + image_studio.replace_selections(project.id, "main", [main.id], path=cfg["db_path"]) + image_studio.replace_selections(project.id, "detail", [detail.id], path=cfg["db_path"]) + return cfg, project, main, detail + + def test_export_partial_selection_outputs_ordered_jpegs(self): + with self.make_temp_dir() as temp_dir: + cfg, project, _main, _detail = self._project_with_assets(temp_dir) + parent = os.path.join(temp_dir, "exports") + os.makedirs(parent) + + result = image_studio_export.export_project_selection( + project.id, + parent, + path=cfg["db_path"], + config=cfg, + ) + + self.assertEqual(1, result.main_count) + self.assertEqual(1, result.detail_count) + self.assertEqual( + ["51100639510_主图_1.jpg", "51100639510_详情图_1.jpg"], + sorted(os.path.basename(item.output_path) for item in result.files), + ) + from PIL import Image + + for item in result.files: + with Image.open(item.output_path) as image: + self.assertEqual("JPEG", image.format) + self.assertEqual("RGB", image.mode) + + self.assert_removed(temp_dir) + + def test_existing_target_requires_explicit_choice_and_timestamp_mode(self): + with self.make_temp_dir() as temp_dir: + cfg, project, _main, _detail = self._project_with_assets(temp_dir) + parent = os.path.join(temp_dir, "exports") + target = os.path.join(parent, project.item_id) + os.makedirs(target) + + with self.assertRaises(image_studio_export.ExportTargetExistsError): + image_studio_export.export_project_selection( + project.id, + parent, + path=cfg["db_path"], + config=cfg, + ) + + result = image_studio_export.export_project_selection( + project.id, + parent, + existing_mode=image_studio_export.EXISTING_TIMESTAMP, + timestamp="20260711_120000", + path=cfg["db_path"], + config=cfg, + ) + self.assertTrue(result.target_dir.endswith("51100639510_20260711_120000")) + self.assertTrue(os.path.isdir(result.target_dir)) + + self.assert_removed(temp_dir) + + def test_overwrite_managed_keeps_user_files(self): + with self.make_temp_dir() as temp_dir: + cfg, project, _main, _detail = self._project_with_assets(temp_dir) + parent = os.path.join(temp_dir, "exports") + target = os.path.join(parent, project.item_id) + os.makedirs(target) + managed = os.path.join(target, "51100639510_主图_9.jpg") + user_file = os.path.join(target, "用户说明.txt") + with open(managed, "w", encoding="utf-8") as fh: + fh.write("old") + with open(user_file, "w", encoding="utf-8") as fh: + fh.write("keep") + + result = image_studio_export.export_project_selection( + project.id, + parent, + existing_mode=image_studio_export.EXISTING_OVERWRITE_MANAGED, + path=cfg["db_path"], + config=cfg, + ) + + self.assertFalse(os.path.exists(managed)) + self.assertTrue(os.path.exists(user_file)) + with open(user_file, encoding="utf-8") as fh: + self.assertEqual("keep", fh.read()) + self.assertEqual(2, len(result.files)) + + self.assert_removed(temp_dir) + + def test_preflight_failure_does_not_create_target(self): + with self.make_temp_dir() as temp_dir: + cfg, project, main, _detail = self._project_with_assets(temp_dir) + os.remove(main.local_path) + parent = os.path.join(temp_dir, "exports") + os.makedirs(parent) + + with self.assertRaisesRegex(image_studio_export.ImageStudioExportError, "缺失"): + image_studio_export.export_project_selection( + project.id, + parent, + path=cfg["db_path"], + config=cfg, + ) + + self.assertFalse(os.path.exists(os.path.join(parent, project.item_id))) + + self.assert_removed(temp_dir) + + +if __name__ == "__main__": + unittest.main()