[~] Refactor

This commit is contained in:
Siarhei Siniak 2024-07-06 15:13:45 +03:00
parent c28d9a19cc
commit 7ade65b678

@ -1,9 +1,13 @@
import logging import logging
import dataclasses
import traceback
import subprocess
import os import os
import sys import sys
import json import json
from typing import ( from typing import (
Any, Any,
Literal,
Optional, Optional,
Iterable, Iterable,
) )
@ -12,6 +16,36 @@ logger = logging.getLogger(__name__)
#logging.getLogger().setLevel(logging.INFO) #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(): def logger_setup():
if len(logger.handlers) == 0: if len(logger.handlers) == 0:
handler = logging.StreamHandler(sys.stderr) handler = logging.StreamHandler(sys.stderr)
@ -89,17 +123,22 @@ def tiktok_videos_meta(links: Iterable[str]) -> Iterable[dict[str, Any]]:
res.append(dict( res.append(dict(
url=o, url=o,
id=int(parts[-1]),
fname='_'.join(parts[-3:]) +'.mp4', fname='_'.join(parts[-3:]) +'.mp4',
)) ))
return res return res
def tiktok_videos_fetch( 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]]: ) -> Iterable[dict[str, Any]]:
import pyktok import pyktok
import tqdm import tqdm
if method is None:
method = 'pyktok'
stats = dict( stats = dict(
saved=0, saved=0,
total=0, total=0,
@ -111,9 +150,19 @@ def tiktok_videos_fetch(
stats['total'] += 1 stats['total'] += 1
if not os.path.exists(o['fname']): if not os.path.exists(o['fname']):
try: try:
if method == 'pyktok':
pyktok.save_tiktok(o['url']) 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 stats['saved'] += 1
except: except:
logger.error(json.dumps(dict(
msg=traceback.format_exc(),
)))
stats['error'] += 1 stats['error'] += 1
else: else:
stats['skipped'] += 1 stats['skipped'] += 1