| 1 | from __future__ import annotations |
| 2 | |
| 3 | from typing import Any, Dict, List, Optional |
| 4 | |
| 5 | from core.user_modes.base_strategy import BaseUserModeStrategy |
| 6 | from utils.logger import setup_logger |
| 7 | |
| 8 | logger = setup_logger("CollectUserModeStrategy") |
| 9 | |
| 10 | |
| 11 | class CollectUserModeStrategy(BaseUserModeStrategy): |
| 12 | mode_name = "collect" |
| 13 | api_method_name = "get_user_collects" |
| 14 | |
| 15 | def __init__(self, downloader, *, collects_id: Optional[str] = None): |
| 16 | """Optional ``collects_id`` constrains the collection to a single |
| 17 | folder. When provided we skip the ``get_user_collects`` enumeration |
| 18 | entirely and only paginate ``get_collect_aweme(collects_id, ...)`` |
| 19 | for that one folder — used by the desktop "我的内容 / 我的收藏" |
| 20 | sub-tab when the user clicks "下载本收藏夹". When ``None`` (CLI |
| 21 | default and historic desktop behaviour), we enumerate every folder |
| 22 | on the account. |
| 23 | """ |
| 24 | super().__init__(downloader) |
| 25 | self._collects_id_filter = (collects_id or "").strip() or None |
| 26 | |
| 27 | async def collect_items(self, sec_uid: str, user_info: Dict[str, Any]) -> List[Dict[str, Any]]: |
| 28 | if self._collects_id_filter: |
| 29 | return await self._collect_single_folder(self._collects_id_filter) |
| 30 | return await self._collect_all_folders(sec_uid) |
| 31 | |
| 32 | async def _collect_single_folder(self, collects_id: str) -> List[Dict[str, Any]]: |
| 33 | """Paginate aweme entries for a single collection folder. |
| 34 | |
| 35 | Mirrors the inner loop of :meth:`_collect_all_folders` but |
| 36 | intentionally avoids :meth:`api_client.get_user_collects` so we |
| 37 | never even read the names of other folders on the account |
| 38 | (Property 4 / R6.4 — single-folder filter does not leak entries |
| 39 | from sibling folders). |
| 40 | """ |
| 41 | fetch_collect_aweme = getattr(self.downloader.api_client, "get_collect_aweme", None) |
| 42 | if not callable(fetch_collect_aweme): |
| 43 | logger.warning("API client missing get_collect_aweme") |
| 44 | return [] |
| 45 | |
| 46 | expanded: List[Dict[str, Any]] = [] |
| 47 | seen_aweme: set[str] = set() |
| 48 | |
| 49 | cursor = 0 |
| 50 | has_more = True |
| 51 | while has_more: |
| 52 | await self.downloader.rate_limiter.acquire() |
| 53 | page_data = await fetch_collect_aweme(str(collects_id), max_cursor=cursor, count=20) |
| 54 | page = self._normalize_page_data(page_data) |
| 55 | page_items = page.get("items", []) |
| 56 | if not page_items: |
| 57 | break |
| 58 | |
| 59 | for item in page_items: |
| 60 | aweme = self._extract_aweme_from_item(item) |
| 61 | if not aweme: |
| 62 | continue |
| 63 | aweme_id = str(aweme.get("aweme_id") or "") |
| 64 | if not aweme_id or aweme_id in seen_aweme: |
| 65 | continue |
| 66 | seen_aweme.add(aweme_id) |
| 67 | expanded.append(aweme) |
| 68 | |
| 69 | has_more = bool(page.get("has_more", False)) |
| 70 | next_cursor = int(page.get("max_cursor", 0) or 0) |
| 71 | if has_more and next_cursor == cursor: |
| 72 | logger.warning("Collect folder %s cursor did not advance", collects_id) |
| 73 | break |
| 74 | cursor = next_cursor |
| 75 | |
| 76 | return expanded |
| 77 | |
| 78 | async def _collect_all_folders(self, sec_uid: str) -> List[Dict[str, Any]]: |
| 79 | """Original behaviour: enumerate every folder on the account and |
| 80 | paginate each one. Kept as a separate method so the filter branch |
| 81 | in :meth:`collect_items` doesn't accidentally invoke |
| 82 | :meth:`api_client.get_user_collects`. |
| 83 | """ |
| 84 | fetch_collect_aweme = getattr(self.downloader.api_client, "get_collect_aweme", None) |
| 85 | fetch_collects = getattr(self.downloader.api_client, self.api_method_name, None) |
| 86 | if not callable(fetch_collects): |
| 87 | logger.warning("API client missing %s", self.api_method_name) |
| 88 | return [] |
| 89 | if not callable(fetch_collect_aweme): |
| 90 | logger.warning("API client missing get_collect_aweme") |
| 91 | return [] |
| 92 | |
| 93 | raw_collects = await self._collect_paged_entries(fetch_collects, sec_uid) |
| 94 | expanded: List[Dict[str, Any]] = [] |
| 95 | seen_aweme: set[str] = set() |
| 96 | |
| 97 | for collect_item in raw_collects: |
| 98 | collects_id = self._extract_collects_id(collect_item) |
| 99 | if not collects_id: |
| 100 | continue |
| 101 | |
| 102 | cursor = 0 |
| 103 | has_more = True |
| 104 | while has_more: |
| 105 | await self.downloader.rate_limiter.acquire() |
| 106 | page_data = await fetch_collect_aweme(str(collects_id), max_cursor=cursor, count=20) |
| 107 | page = self._normalize_page_data(page_data) |
| 108 | page_items = page.get("items", []) |
| 109 | if not page_items: |
| 110 | break |
| 111 | |
| 112 | for item in page_items: |
| 113 | aweme = self._extract_aweme_from_item(item) |
| 114 | if not aweme: |
| 115 | continue |
| 116 | aweme_id = str(aweme.get("aweme_id") or "") |
| 117 | if not aweme_id or aweme_id in seen_aweme: |
| 118 | continue |
| 119 | seen_aweme.add(aweme_id) |
| 120 | expanded.append(aweme) |
| 121 | |
| 122 | has_more = bool(page.get("has_more", False)) |
| 123 | next_cursor = int(page.get("max_cursor", 0) or 0) |
| 124 | if has_more and next_cursor == cursor: |
| 125 | logger.warning("Collect folder %s cursor did not advance", collects_id) |
| 126 | break |
| 127 | cursor = next_cursor |
| 128 | |
| 129 | return expanded |
| 130 | |
| 131 | @staticmethod |
| 132 | def _extract_collects_id(item: Any) -> str: |
| 133 | if not isinstance(item, dict): |
| 134 | return "" |
| 135 | return str( |
| 136 | item.get("collects_id") |
| 137 | or item.get("collects_id_str") |
| 138 | or item.get("id") |
| 139 | or ((item.get("collects_info") or {}).get("collects_id")) |
| 140 | or ((item.get("collects_info") or {}).get("collects_id_str")) |
| 141 | or "" |
| 142 | ) |
| 143 |