feat(product-suite): make image pulls cancellable

app/image_studio_images.py

@@ -32,6 +32,10 @@ class ImageStudioImageError(RuntimeError):
     """Raised when a remote image cannot be safely loaded or saved."""
 
 
+class ImageStudioImageCancelled(ImageStudioImageError):
+    """Raised when an original image download is cancelled safely."""
+
+
 @dataclass(frozen=True)
 class RemoteImage:
     url: str
@@ -127,10 +131,14 @@ def download_remote_image(
     max_bytes=ORIGINAL_MAX_BYTES,
     timeout=(DEFAULT_CONNECT_TIMEOUT_SECONDS, DEFAULT_READ_TIMEOUT_SECONDS),
     session=None,
+    should_stop=None,
 ) -> RemoteImage:
     """Download remote image bytes with SSRF, timeout and size guards."""
 
     url = str(url or "").strip()
+    should_stop = should_stop or (lambda: False)
+    if should_stop():
+        raise ImageStudioImageCancelled("用户停止下载商品原图")
     _assert_public_http_url(url)
     client = session or _session()
     try:
@@ -142,38 +150,55 @@ def download_remote_image(
         )
     except requests.exceptions.RequestException as exc:
         raise ImageStudioImageError(f"远程图片下载失败：{exc}") from exc
-    status = getattr(response, "status_code", 200)
-    if status >= 400:
-        raise ImageStudioImageError(f"远程图片下载失败：HTTP {status}")
-    final_url = str(getattr(response, "url", "") or url)
-    _assert_public_http_url(final_url)
-    content_type = _validate_content_type(getattr(response, "headers", {}).get("Content-Type", ""))
-    content_length = getattr(response, "headers", {}).get("Content-Length")
-    if content_length:
-        try:
-            if int(content_length) > int(max_bytes):
+    try:
+        if should_stop():
+            raise ImageStudioImageCancelled("用户停止下载商品原图")
+        status = getattr(response, "status_code", 200)
+        if status >= 400:
+            raise ImageStudioImageError(f"远程图片下载失败：HTTP {status}")
+        final_url = str(getattr(response, "url", "") or url)
+        _assert_public_http_url(final_url)
+        content_type = _validate_content_type(
+            getattr(response, "headers", {}).get("Content-Type", "")
+        )
+        content_length = getattr(response, "headers", {}).get("Content-Length")
+        if content_length:
+            try:
+                if int(content_length) > int(max_bytes):
+                    raise ImageStudioImageError("远程图片超过大小上限")
+            except ValueError:
+                pass
+        chunks = []
+        total = 0
+        iterator = (
+            response.iter_content(chunk_size=65536)
+            if hasattr(response, "iter_content")
+            else [response.content]
+        )
+        for chunk in iterator:
+            if should_stop():
+                raise ImageStudioImageCancelled("用户停止下载商品原图")
+            if not chunk:
+                continue
+            total += len(chunk)
+            if total > int(max_bytes):
                 raise ImageStudioImageError("远程图片超过大小上限")
-        except ValueError:
-            pass
-    chunks = []
-    total = 0
-    iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content]
-    for chunk in iterator:
-        if not chunk:
-            continue
-        total += len(chunk)
-        if total > int(max_bytes):
-            raise ImageStudioImageError("远程图片超过大小上限")
-        chunks.append(chunk)
-    content = b"".join(chunks)
-    _image_info(content)
-    return RemoteImage(
-        url=url,
-        content=content,
-        content_type=content_type,
-        final_url=final_url,
-        redirected=final_url != url,
-    )
+            chunks.append(chunk)
+        if should_stop():
+            raise ImageStudioImageCancelled("用户停止下载商品原图")
+        content = b"".join(chunks)
+        _image_info(content)
+        return RemoteImage(
+            url=url,
+            content=content,
+            content_type=content_type,
+            final_url=final_url,
+            redirected=final_url != url,
+        )
+    finally:
+        close = getattr(response, "close", None)
+        if callable(close):
+            close()
 
 
 def _image_info(image_bytes):
@@ -509,10 +534,21 @@ def restore_trashed_asset(record, *, path=None, config=None):
         raise ImageStudioImageError(f"撤销删除生成图片失败：{exc}") from exc
 
 
-def download_original_asset(asset_id, *, path=None, config=None, image_root=None, session=None):
+def download_original_asset(
+    asset_id,
+    *,
+    path=None,
+    config=None,
+    image_root=None,
+    session=None,
+    should_stop=None,
+):
     """Download one Shopee original image into originals/ and mark its asset available."""
 
     cfg = appconfig.load_config() if config is None else config
+    should_stop = should_stop or (lambda: False)
+    if should_stop():
+        raise ImageStudioImageCancelled("用户停止下载商品原图")
     database_path = path or appconfig.db_path(cfg)
     asset = image_studio.get_asset(asset_id, path=database_path)
     if asset is None:
@@ -525,15 +561,26 @@ def download_original_asset(asset_id, *, path=None, config=None, image_root=None
     project = image_studio.get_project(asset.project_id, path=database_path)
     if project is None:
         raise ImageStudioImageError("原图资产所属项目不存在")
-    remote = download_remote_image(asset.remote_url, max_bytes=ORIGINAL_MAX_BYTES, session=session)
+    remote = download_remote_image(
+        asset.remote_url,
+        max_bytes=ORIGINAL_MAX_BYTES,
+        session=session,
+        should_stop=should_stop,
+    )
+    if should_stop():
+        raise ImageStudioImageCancelled("用户停止下载商品原图")
     info = _image_info(remote.content)
     directory, stem = _original_file_path(project, asset, image_root or appconfig.image_dir(cfg))
     os.makedirs(directory, exist_ok=True)
     final_path = os.path.join(directory, stem + _extension_for_format(info["format"]))
     temp_path = final_path + ".tmp-" + uuid.uuid4().hex
     try:
+        if should_stop():
+            raise ImageStudioImageCancelled("用户停止下载商品原图")
         with open(temp_path, "wb") as fh:
             fh.write(remote.content)
+        if should_stop():
+            raise ImageStudioImageCancelled("用户停止下载商品原图")
         with open(temp_path, "rb") as fh:
             _image_info(fh.read())
         os.replace(temp_path, final_path)