toolbox/live_info_collect/manager.py

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

from toolbox.live_info_collect.douyin import Douyin
from project_settings import project_path

toolbox_logger = logging.getLogger("toolbox")


class LiveInfoCollectManager(object):
    def __init__(self, live_info_collect_file: str):
        self.live_info_collect_file = live_info_collect_file

        # state
        self.coro_task_dict = dict()

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

        for task in tasks:
            room_platform = task["platform"]
            check_interval = task["check_interval"]
            file_output_file = project_path / task["file_output_file"]

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

            record_task = Douyin(
                platform=room_platform,
                check_interval=check_interval,
                file_output_file=file_output_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)

    live_info_collect_tasks_file = project_path / "data/live_info_collect.json"

    manager = LiveInfoCollectManager(live_info_collect_tasks_file)
    await manager.run()
    return


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