mirror of
https://github.com/imgyh/tiktok.git
synced 2025-04-06 08:15:46 +08:00
feat(tiktok): 使用rich进度条
This commit is contained in:
imgyh
parent
ebb60c8ee7
commit
0fb3739b47
125
TikTok.py
125
TikTok.py
@ -20,7 +20,19 @@ import time
|
||||
import os
|
||||
import copy
|
||||
from concurrent.futures import ThreadPoolExecutor, wait, ALL_COMPLETED
|
||||
from tqdm import tqdm
|
||||
from functools import partial
|
||||
from urllib.request import urlopen
|
||||
import signal
|
||||
from threading import Event
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
DownloadColumn,
|
||||
Progress,
|
||||
TaskID,
|
||||
TextColumn,
|
||||
TimeRemainingColumn,
|
||||
TransferSpeedColumn
|
||||
)
|
||||
|
||||
from TikTokUtils import Utils
|
||||
from TikTokUrls import Urls
|
||||
@ -41,6 +53,21 @@ class TikTok(object):
|
||||
# 用于设置重复请求某个接口的最大时间
|
||||
self.timeout = 10
|
||||
|
||||
# rich 进度条
|
||||
self.progress = Progress(
|
||||
TextColumn("[bold blue]{task.fields[filename]}", justify="left"),
|
||||
BarColumn(bar_width=20),
|
||||
"[progress.percentage]{task.percentage:>3.1f}%",
|
||||
"•",
|
||||
DownloadColumn(),
|
||||
"•",
|
||||
TransferSpeedColumn(),
|
||||
"•",
|
||||
TimeRemainingColumn(),
|
||||
)
|
||||
self.done_event = Event()
|
||||
signal.signal(signal.SIGINT, self.handle_sigint)
|
||||
|
||||
|
||||
# 从分享链接中提取网址
|
||||
def getShareLink(self, string):
|
||||
@ -513,32 +540,31 @@ class TikTok(object):
|
||||
|
||||
return awemeList
|
||||
|
||||
# 来自 https://blog.csdn.net/weixin_43347550/article/details/105248223
|
||||
def progressBarDownload(self, url, filepath,desc):
|
||||
response = requests.get(url, stream=True, headers=self.headers)
|
||||
chunk_size = 1024 # 每次下载的数据大小
|
||||
content_size = int(response.headers['content-length']) # 下载文件总大小
|
||||
try:
|
||||
if response.status_code == 200: # 判断是否响应成功
|
||||
# print('[开始下载]:文件大小:{size:.2f} MB'.format(
|
||||
# size=content_size / chunk_size / 1024)) # 开始下载，显示下载文件大小
|
||||
with open(filepath, 'wb') as file, tqdm(total=content_size,
|
||||
unit="iB",
|
||||
desc=desc,
|
||||
unit_scale=True,
|
||||
unit_divisor=1024,
|
||||
# https://github.com/textualize/rich/blob/master/examples/downloader.py
|
||||
def handle_sigint(self, signum, frame):
|
||||
self.done_event.set()
|
||||
|
||||
) as bar: # 显示进度条
|
||||
for data in response.iter_content(chunk_size=chunk_size):
|
||||
size = file.write(data)
|
||||
bar.update(size)
|
||||
def copy_url(self, task_id: TaskID, url: str, path: str) -> None:
|
||||
"""Copy data from a url to a local file."""
|
||||
# self.progress.console.log(f"Requesting {url}")
|
||||
response = urlopen(url)
|
||||
try:
|
||||
# This will break if the response doesn't contain content length
|
||||
self.progress.update(task_id, total=int(response.info()["Content-length"]))
|
||||
with open(path, "wb") as dest_file:
|
||||
self.progress.start_task(task_id)
|
||||
for data in iter(partial(response.read, 32768), b""):
|
||||
dest_file.write(data)
|
||||
self.progress.update(task_id, advance=len(data))
|
||||
if self.done_event.is_set():
|
||||
return
|
||||
except Exception as e:
|
||||
# 下载异常 删除原来下载的文件, 可能未下成功
|
||||
if os.path.exists(filepath):
|
||||
os.remove(filepath)
|
||||
if os.path.exists(path):
|
||||
os.remove(path)
|
||||
print("[ 错误 ]:下载出错\r")
|
||||
|
||||
def awemeDownload(self, awemeDict: dict, music=True, cover=True, avatar=True, resjson=True, savePath=os.getcwd(), usingThread = True):
|
||||
def awemeDownload(self, awemeDict: dict, music=True, cover=True, avatar=True, resjson=True, savePath=os.getcwd()):
|
||||
if awemeDict is None:
|
||||
return
|
||||
if not os.path.exists(savePath):
|
||||
@ -574,10 +600,8 @@ class TikTok(object):
|
||||
try:
|
||||
url = awemeDict["video"]["play_addr"]["url_list"]
|
||||
if url != "":
|
||||
if usingThread:
|
||||
self.alltask.append(self.tpool.submit(self.progressBarDownload, url, video_path, "[ 视频 ]:" + desc))
|
||||
else:
|
||||
self.progressBarDownload(url, video_path, "[ 视频 ]:" + desc)
|
||||
task_id = self.progress.add_task("download", filename="[ 视频 ]:" + desc, start=False)
|
||||
self.alltask.append(self.pool.submit(self.copy_url, task_id, url, video_path))
|
||||
except Exception as e:
|
||||
print("[ 警告 ]:视频下载失败,请重试...\r\n")
|
||||
|
||||
@ -593,10 +617,8 @@ class TikTok(object):
|
||||
try:
|
||||
url = image["url_list"][0]
|
||||
if url != "":
|
||||
if usingThread:
|
||||
self.alltask.append(self.tpool.submit(self.progressBarDownload, url, image_path, "[ 图集 ]:" + desc))
|
||||
else:
|
||||
self.progressBarDownload(url, image_path, "[ 图集 ]:" + desc)
|
||||
task_id = self.progress.add_task("download", filename="[ 图集 ]:" + desc, start=False)
|
||||
self.alltask.append(self.pool.submit(self.copy_url, task_id, url, image_path))
|
||||
except Exception as e:
|
||||
print("[ 警告 ]:图片下载失败,请重试...\r\n")
|
||||
|
||||
@ -613,12 +635,9 @@ class TikTok(object):
|
||||
try:
|
||||
url = awemeDict["music"]["play_url"]["url_list"][0]
|
||||
if url != "":
|
||||
if usingThread:
|
||||
self.alltask.append(self.tpool.submit(self.progressBarDownload, url, music_path, "[ 原声 ]:" + desc))
|
||||
else:
|
||||
self.progressBarDownload(url, music_path, "[ 原声 ]:" + desc)
|
||||
task_id = self.progress.add_task("download", filename="[ 原声 ]:" + desc, start=False)
|
||||
self.alltask.append(self.pool.submit(self.copy_url, task_id, url, music_path))
|
||||
except Exception as e:
|
||||
# print(e)
|
||||
print("[ 警告 ]:音乐(原声)下载失败,请重试...\r\n")
|
||||
|
||||
# 下载 cover
|
||||
@ -633,12 +652,9 @@ class TikTok(object):
|
||||
try:
|
||||
url = awemeDict["video"]["cover_original_scale"]["url_list"][0]
|
||||
if url != "":
|
||||
if usingThread:
|
||||
self.alltask.append(self.tpool.submit(self.progressBarDownload, url, cover_path, "[ 封面 ]:" + desc))
|
||||
else:
|
||||
self.progressBarDownload(url, cover_path, "[ 封面 ]:" + desc)
|
||||
task_id = self.progress.add_task("download", filename="[ 封面 ]:" + desc, start=False)
|
||||
self.alltask.append(self.pool.submit(self.copy_url, task_id, url, cover_path))
|
||||
except Exception as e:
|
||||
# print(e)
|
||||
print("[ 警告 ]:cover下载失败,请重试...\r\n")
|
||||
|
||||
# 下载 avatar
|
||||
@ -653,12 +669,9 @@ class TikTok(object):
|
||||
try:
|
||||
url = awemeDict["author"]["avatar"]["url_list"][0]
|
||||
if url != "":
|
||||
if usingThread:
|
||||
self.alltask.append(self.tpool.submit(self.progressBarDownload, url, avatar_path, "[ 头像 ]:" + desc))
|
||||
else:
|
||||
self.progressBarDownload(url, avatar_path, "[ 头像 ]:" + desc)
|
||||
task_id = self.progress.add_task("download", filename="[ 头像 ]:" + desc, start=False)
|
||||
self.alltask.append(self.pool.submit(self.copy_url, task_id, url, avatar_path))
|
||||
except Exception as e:
|
||||
# print(e)
|
||||
print("[ 警告 ]:avatar下载失败,请重试...\r\n")
|
||||
except Exception as e:
|
||||
print("[ 错误 ]:下载作品时出错\r\n")
|
||||
@ -669,14 +682,26 @@ class TikTok(object):
|
||||
if not os.path.exists(savePath):
|
||||
os.mkdir(savePath)
|
||||
|
||||
self.tpool = ThreadPoolExecutor(thread)
|
||||
self.alltask = []
|
||||
|
||||
start = time.time() # 开始时间
|
||||
for aweme in awemeList:
|
||||
# print("[ 提示 ]:正在下载 [%s] 的作品 %s/%s\r\n"
|
||||
# % (aweme["author"]["nickname"], str(ind + 1), len(awemeList)))
|
||||
self.awemeDownload(awemeDict=aweme, music=music, cover=cover, avatar=avatar, resjson=resjson, savePath=savePath,usingThread=True)
|
||||
with self.progress:
|
||||
with ThreadPoolExecutor(max_workers=thread) as self.pool:
|
||||
self.progress.console.log("请耐心等待下载完成(终端尺寸越长显示的进度条越多)...")
|
||||
for aweme in awemeList:
|
||||
self.awemeDownload(awemeDict=aweme, music=music, cover=cover, avatar=avatar, resjson=resjson, savePath=savePath)
|
||||
# time.sleep(0.5)
|
||||
wait(self.alltask, return_when=ALL_COMPLETED)
|
||||
# 清除上一步的进度条
|
||||
for taskid in self.progress.task_ids:
|
||||
self.progress.remove_task(taskid)
|
||||
# 下载上一步失败的
|
||||
with self.progress:
|
||||
with ThreadPoolExecutor(max_workers=thread) as self.pool:
|
||||
self.progress.console.log("正在重试下载失败的内容...")
|
||||
for aweme in awemeList:
|
||||
self.awemeDownload(awemeDict=aweme, music=music, cover=cover, avatar=avatar, resjson=resjson, savePath=savePath)
|
||||
# time.sleep(0.5)
|
||||
|
||||
wait(self.alltask, return_when=ALL_COMPLETED)
|
||||
end = time.time() # 结束时间
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user