| 1 | from __future__ import annotations |
| 2 | |
| 3 | from typing import Any, Dict, List, Optional |
| 4 | |
| 5 | from core.downloader_base import BaseDownloader, DownloadResult |
| 6 | from core.user_modes.base_strategy import BaseUserModeStrategy |
| 7 | from utils.logger import setup_logger |
| 8 | |
| 9 | logger = setup_logger("MixDownloader") |
| 10 | |
| 11 | |
| 12 | class MixDownloader(BaseDownloader): |
| 13 | async def download(self, parsed_url: Dict[str, Any]) -> DownloadResult: |
| 14 | result = DownloadResult() |
| 15 | |
| 16 | mix_id = parsed_url.get("mix_id") |
| 17 | if not mix_id: |
| 18 | logger.error("No mix_id found in parsed URL") |
| 19 | return result |
| 20 | |
| 21 | aweme_list = await self._collect_mix_aweme_list(str(mix_id)) |
| 22 | |
| 23 | result.total = len(aweme_list) |
| 24 | self._progress_set_item_total(result.total, "合集作品待下载") |
| 25 | self._progress_update_step("下载合集", f"mix_id={mix_id}，待处理 {result.total} 条") |
| 26 | |
| 27 | mix_detail = await self._get_mix_detail(str(mix_id)) |
| 28 | author_name = ( |
| 29 | (mix_detail.get("author") or {}).get("nickname") |
| 30 | if isinstance(mix_detail, dict) |
| 31 | else None |
| 32 | ) or "mix" |
| 33 | |
| 34 | async def _process_aweme(item: Dict[str, Any]): |
| 35 | aweme_id = item.get("aweme_id") |
| 36 | if not aweme_id: |
| 37 | self._progress_advance_item("failed", "missing_aweme_id") |
| 38 | return {"status": "failed", "aweme_id": None} |
| 39 | |
| 40 | if not await self._should_download(str(aweme_id)): |
| 41 | self._progress_advance_item("skipped", str(aweme_id)) |
| 42 | return {"status": "skipped", "aweme_id": aweme_id} |
| 43 | |
| 44 | success = await self._download_aweme_assets(item, author_name, mode="mix") |
| 45 | status = "success" if success else "failed" |
| 46 | self._progress_advance_item(status, str(aweme_id)) |
| 47 | return {"status": status, "aweme_id": aweme_id} |
| 48 | |
| 49 | download_results = await self.queue_manager.download_batch(_process_aweme, aweme_list) |
| 50 | for entry in download_results: |
| 51 | status = entry.get("status") if isinstance(entry, dict) else None |
| 52 | if status == "success": |
| 53 | result.success += 1 |
| 54 | elif status == "skipped": |
| 55 | result.skipped += 1 |
| 56 | else: |
| 57 | result.failed += 1 |
| 58 | return result |
| 59 | |
| 60 | async def _collect_mix_aweme_list(self, mix_id: str) -> List[Dict[str, Any]]: |
| 61 | fetch_mix_aweme = getattr(self.api_client, "get_mix_aweme", None) |
| 62 | if not callable(fetch_mix_aweme): |
| 63 | logger.error("API client has no get_mix_aweme implementation") |
| 64 | return [] |
| 65 | |
| 66 | aweme_list: List[Dict[str, Any]] = [] |
| 67 | has_more = True |
| 68 | cursor = 0 |
| 69 | number_limit = int(self.config.get("number", {}).get("mix", 0) or 0) |
| 70 | |
| 71 | while has_more: |
| 72 | await self.rate_limiter.acquire() |
| 73 | raw_page = await fetch_mix_aweme(mix_id, cursor=cursor, count=20) |
| 74 | page = BaseUserModeStrategy._normalize_page_data(raw_page) |
| 75 | items = page.get("items", []) |
| 76 | if not items: |
| 77 | break |
| 78 | |
| 79 | for item in items: |
| 80 | aweme = self._extract_aweme_from_item(item) |
| 81 | if aweme: |
| 82 | aweme_list.append(aweme) |
| 83 | |
| 84 | if number_limit > 0 and len(aweme_list) >= number_limit: |
| 85 | aweme_list = aweme_list[:number_limit] |
| 86 | break |
| 87 | |
| 88 | has_more = bool(page.get("has_more", False)) |
| 89 | next_cursor = int(page.get("max_cursor", 0) or 0) |
| 90 | if has_more and next_cursor == cursor: |
| 91 | logger.warning( |
| 92 | "Mix pagination cursor did not advance (%s), stop to avoid loop", |
| 93 | cursor, |
| 94 | ) |
| 95 | break |
| 96 | cursor = next_cursor |
| 97 | |
| 98 | return aweme_list |
| 99 | |
| 100 | async def _get_mix_detail(self, mix_id: str) -> Optional[Dict[str, Any]]: |
| 101 | getter = getattr(self.api_client, "get_mix_detail", None) |
| 102 | if not callable(getter): |
| 103 | return None |
| 104 | try: |
| 105 | return await getter(mix_id) |
| 106 | except Exception as exc: |
| 107 | logger.warning("Get mix detail failed: %s", exc) |
| 108 | return None |
| 109 | |
| 110 | @staticmethod |
| 111 | def _extract_aweme_from_item(item: Any) -> Optional[Dict[str, Any]]: |
| 112 | if not isinstance(item, dict): |
| 113 | return None |
| 114 | if item.get("aweme_id"): |
| 115 | return item |
| 116 | for key in ("aweme", "aweme_info", "aweme_detail"): |
| 117 | value = item.get(key) |
| 118 | if isinstance(value, dict) and value.get("aweme_id"): |
| 119 | return value |
| 120 | return None |
| 121 |