toolbox/live_recorder/manager.py

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

from streamlink_cli.output import FileOutput

from toolbox.live_recorder.douyin import DouyinLiveRecorder
from project_settings import project_path

toolbox_logger = logging.getLogger("toolbox")


class LiveRecorderManager(object):
    def __init__(self,
                 live_recorder_tasks_file: str,
                 ):
        self.live_recorder_tasks_file = live_recorder_tasks_file

        # state
        self.coro_task_dict = dict()

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

        for task in tasks:
            room_platform = task["platform"]
            room_id = task["room_id"]
            room_name = task["room_name"]
            check_interval = task["check_interval"]
            file_output_fmt = task["file_output_fmt"]
            file_output_dir = project_path / task["file_output_dir"]
            file_info_file = project_path / task["file_info_file"]

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

            record_task = DouyinLiveRecorder(
                platform=room_platform,
                room_id=room_id,
                room_name=room_name,
                check_interval=check_interval,
                file_output_fmt=file_output_fmt,
                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)

    live_recorder_tasks_file = project_path / "data/live_recorder_tasks.json"

    manager = LiveRecorderManager(live_recorder_tasks_file)
    await manager.run()
    return


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