diff --git a/python/tasks/tiktok/__init__.py b/python/tasks/tiktok/__init__.py index 25052c6..914dec8 100644 --- a/python/tasks/tiktok/__init__.py +++ b/python/tasks/tiktok/__init__.py @@ -66,6 +66,7 @@ async def tiktok_videos_links_get( query: Optional[str]=None, screenshot_path: Optional[str]=None, max_time: Optional[int | float]=None, + max_links: Optional[int]=None, ) -> Iterable[str]: import datetime import TikTokApi @@ -73,6 +74,9 @@ async def tiktok_videos_links_get( import asyncio import re + if max_links is None: + max_links = 100 + if max_time is None: max_time = 10 @@ -91,7 +95,8 @@ async def tiktok_videos_links_get( path=screenshot_path, ) - links = set() + links = list() + links_set = set() started_at = datetime.datetime.now() @@ -104,7 +109,9 @@ async def tiktok_videos_links_get( old_size = len(links) for o in new_links: - links.add(o) + if not o in links_set: + links_set.add(o) + links.append(o) await session.page.mouse.wheel(0, 100) @@ -115,6 +122,9 @@ async def tiktok_videos_links_get( if elapsed > max_time: break; + if len(links_set) > max_links: + break + if old_size < len(links): logger.info(json.dumps(dict( total=len(links), @@ -122,7 +132,7 @@ async def tiktok_videos_links_get( scroll_y=await session.page.evaluate('window.scrollY'), ))) - return links + return list(links)[:max_links] def tiktok_videos_meta(links: Iterable[str]) -> Iterable[dict[str, Any]]: res = []