diff --git a/python/tasks/tiktok/__init__.py b/python/tasks/tiktok/__init__.py index 9e9035b..63ca13e 100644 --- a/python/tasks/tiktok/__init__.py +++ b/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