From 45e289cea7cdbbeea1501df88258dacea95de460 Mon Sep 17 00:00:00 2001
From: Siarhei Siniak <siarheisiniak@gmail.com>
Date: Sat, 6 Jul 2024 17:55:04 +0300
Subject: [PATCH] [~] Refactor
---
python/tasks/tiktok/__init__.py | 16 +++++++++++++---
1 file changed, 13 insertions(+), 3 deletions(-)
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 = []