Merge pull request #6 from captaincolonelfox/dev

Release 1.4.0

README.md

@@ -1,4 +1,4 @@
-[![TeleTok](./img/logo.jpg?raw=true)](https://t.me/shieldy_bot)
+[![TeleTok](./img/logo.jpg?raw=true)](https://t.me/TeleTockerBot)
 
 # [TeleTok](https://t.me/TeleTockerBot): Telegram bot for TikTok (and Likee)
 

bot/api/__init__.py

@@ -1,3 +1,4 @@
+import asyncio
 import re
 from abc import ABC, abstractmethod
 from typing import Any, List, Optional
@@ -46,17 +47,19 @@ async def handle_message(self, message: Message) -> List[VideoData]:
             sentry_sdk.capture_exception(ex)
             return []
 
-    async def download_video(self, url: str) -> VideoData:
-        page = await self.client.get(url)
-        soup = BeautifulSoup(page.content, 'html.parser')
-        if data := soup(text=re.compile(self.regexp_key)):
-            for script in data:
-                if link := self._parse_data(script):
-                    video = await self.client.get(link, headers=self.headers)
-                    video.raise_for_status()
-                    return VideoData(link, video.content)
-                return VideoData()
+    async def download_video(self, url: str, retries: int = 2) -> VideoData:
+        for _ in range(retries):
+            page = await self.client.get(url)
+            soup = BeautifulSoup(page.content, 'html.parser')
+            if data := soup(text=re.compile(self.regexp_key)):
+                for script in data:
+                    if link := self._parse_data(script):
+                        if video := await self.client.get(link, headers=self.headers):
+                            video.raise_for_status()
+                            return VideoData(link, video.content)
+            await asyncio.sleep(0.5)
+        return VideoData()
 
     @abstractmethod
-    def _parse_data(self, script: str) -> str:
+    def _parse_data(self, script: str) -> Optional[str]:
         pass

bot/api/likee.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any
+from typing import Any, Optional
 
 from bot.api import API
 
@@ -18,6 +18,6 @@ def links(self):
     def regexp_key(self) -> str:
         return r'"video_url":'
 
-    def _parse_data(self, script: str) -> str:
+    def _parse_data(self, script: str) -> Optional[str]:
         js = json.loads(script.split('window.data = ')[-1][:-1])
-        return js['video_url']
+        return js.get('video_url')

bot/api/tiktok.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any
+from typing import Any, Optional
 
 from bot.api import API
 
@@ -18,6 +18,11 @@ def links(self):
     def regexp_key(self) -> str:
         return r'"downloadAddr":"'
 
-    def _parse_data(self, script: str) -> str:
+    def _parse_data(self, script: str) -> Optional[str]:
         js = json.loads(script)
-        return js['props']['pageProps']['itemInfo']['itemStruct']['video']['downloadAddr']
+        return js.get('props', {})\
+                 .get('pageProps', {})\
+                 .get('itemInfo', {})\
+                 .get('itemStruct', {})\
+                 .get('video', {})\
+                 .get('downloadAddr')