toolbox/video_downloader/manager.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import asyncio
import json
import logging

from toolbox.video_downloader.douyin import DouyinDownloader
from project_settings import project_path

toolbox_logger = logging.getLogger("toolbox")


class VideoDownloaderManager(object):
    def __init__(self,
                 video_download_tasks_file: str,
                 ):
        self.video_download_tasks_file = video_download_tasks_file

        # state
        self.coro_task_dict = dict()

    def get_init_tasks(self):
        with open(self.video_download_tasks_file, "r", encoding="utf-8") as f:
            tasks = json.load(f)

        for task in tasks:
            platform = task["platform"]
            user_name = task["user_name"]
            sec_user_id = task["sec_user_id"]
            min_date = task.get("min_date")
            check_interval = task["check_interval"]
            file_output_dir = project_path / task["file_output_dir"]
            file_info_file = project_path / task["file_info_file"]

            key = f"{platform}_{sec_user_id}"
            if key in self.coro_task_dict.keys():
                continue

            record_task = DouyinDownloader(
                platform=platform,
                user_name=user_name,
                sec_user_id=sec_user_id,
                min_date=min_date,
                check_interval=check_interval,
                file_output_dir=file_output_dir,
                file_info_file=file_info_file,
            )
            self.coro_task_dict[key] = record_task.start()

        future_tasks = [asyncio.create_task(task) for task in self.coro_task_dict.values()]
        return future_tasks

    async def run(self):
        future_tasks = self.get_init_tasks()
        await asyncio.wait(future_tasks)


async def main():
    import log
    from project_settings import project_path, log_directory

    log.setup_size_rotating(log_directory=log_directory)

    video_download_tasks_file = project_path / "data/video_download_tasks.json"

    manager = VideoDownloaderManager(video_download_tasks_file)
    await manager.run()
    return


if __name__ == "__main__":
    asyncio.run(main())