[~] Refactor

python/tasks/tiktok/__init__.py

@@ -1,9 +1,13 @@
 import logging
+import dataclasses
+import traceback
+import subprocess
 import os
 import sys
 import json
 from typing import (
     Any,
+    Literal,
     Optional,
     Iterable,
 )
@@ -12,6 +16,36 @@ logger = logging.getLogger(__name__)
 
 #logging.getLogger().setLevel(logging.INFO)
 
+class tiktok_config_t:
+    @dataclasses.dataclass
+    class res_t:
+        project_root: str=''
+        cache: str=''
+        videos: str=''
+
+def tiktok_config() -> tiktok_config_t.res_t:
+    res = tiktok_config_t.res_t(
+        project_root=os.path.abspath(
+            os.path.join(
+                __file__,
+                '..', '..', '..',
+            ),
+        ),
+    )
+
+    res.cache = os.path.join(
+        res.project_root,
+        'tmp/cache/tiktok',
+    )
+    res.videos = os.path.join(
+        res.cache,
+        'videos',
+    )
+
+    os.makedirs(res.videos, exist_ok=True)
+
+    return res
+
 def logger_setup():
     if len(logger.handlers) == 0:
         handler = logging.StreamHandler(sys.stderr)
@@ -89,17 +123,22 @@ def tiktok_videos_meta(links: Iterable[str]) -> Iterable[dict[str, Any]]:
 
         res.append(dict(
             url=o,
+            id=int(parts[-1]),
             fname='_'.join(parts[-3:]) +'.mp4',
         ))
 
     return res
 
 def tiktok_videos_fetch(
-    meta: Iterable[dict[str, Any]]
+    meta: Iterable[dict[str, Any]],
+    method: Optional[Literal['pyktok', 'tikcdn.io']]=None,
 ) -> Iterable[dict[str, Any]]:
     import pyktok
     import tqdm
 
+    if method is None:
+        method = 'pyktok'
+
     stats = dict(
         saved=0,
         total=0,
@@ -111,9 +150,19 @@ def tiktok_videos_fetch(
         stats['total'] += 1
         if not os.path.exists(o['fname']):
             try:
-                pyktok.save_tiktok(o['url'])
+                if method == 'pyktok':
+                    pyktok.save_tiktok(o['url'])
+                elif method == 'tikcdn.io':
+                    subprocess.check_call([
+                        'curl',
+                        'https://tikcdn.io/ssstik/%d' % o['id'],
+                        '-o', o['fname'],
+                    ])
                 stats['saved'] += 1
             except:
+                logger.error(json.dumps(dict(
+                    msg=traceback.format_exc(),
+                )))
                 stats['error'] += 1
         else:
             stats['skipped'] += 1