python3 视频分段下载

视频分段下载器

#!/usr/bin/env python
# -*- encoding: utf-8 -*-

"""
@author      : sirius
@file        : block_downloader.py
@time        : 2023/1/17
@description :
"""
import logging
import os
import time
import requests
from logging import handlers, Formatter
from concurrent.futures import ThreadPoolExecutor
from queue import Queue

logger = logging.getLogger(__name__)

class BlockDownLoader():
    """
        视频分段下载器
            1. 获取视频总长度 requests获取header['content-length']字段
            2. 根据块大小，对视频长度进行分割
            3. 把分割结果推到下载内部任务队列
            4. 对任务队列中的任务进行并发下载，并写入文件
            5. 对文件进行校验，合并本地文件片段，删除分割片段
    """

    def __init__(self, url, video_key, block_size=1024 * 1024 * 5, remove_block_file=True, max_threads=5, dir_prefix=None):
        """
        :param url: 视频可下载链接这里主要处理可直接下载的云链接
        :param video_key: 云存储的key,也作为视频文件下载生成路径目录的参考，主要目的是为了防止下载的文件相互覆盖
        :param block_size: 分块下载块的大小 默认5M
        :param remove_block_file: 是否清楚分块下载的临时文件 默认是
        :param max_threads: 最大线程数 默认5
        :param dir_prefix: 下载目录文件前缀目录
        """
        self._init_log()
        curr_dirname = os.path.abspath(os.path.dirname(__file__))
        output_path, ext = os.path.splitext(video_key)
        dir_path, file_name = os.path.split(video_key)

        self.url = url
        if dir_prefix:
            self.output_path = f"{dir_prefix}/{output_path}"
        else:
            self.output_path = f"{curr_dirname}/{output_path}"
        self.file_name = file_name
        self.block_size = block_size
        self.max_threads = max_threads
        self.remove_block_file = remove_block_file
        self.q_file_section = Queue()

        self.file_path = f"{self.output_path}/{self.file_name}"
        self.section_count = 0
        self.video_size = 0
        logger.info(f"curr_path={curr_dirname}")
        logger.info(f"dir_path={self.output_path}")
        logger.info(f"file_name={self.file_name}")
        logger.info(f"file_path={self.file_path}")

    def _ensure_dir(self, local_path):
        dirname = os.path.dirname(local_path)
        if not os.path.isdir(dirname):
            os.makedirs(dirname)

    def _init_log(self):
        CURRENT_PATH = os.path.dirname(os.path.abspath(__file__))
        ROOT_PATH = os.path.join(CURRENT_PATH, os.pardir)
        logger_dir = os.path.join(ROOT_PATH, 'logs')
        self._ensure_dir(logger_dir)

        logger = logging.getLogger()
        logger_path = os.path.join(logger_dir, 'downloader.log')
        format = "%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]"
        file_handler = logging.handlers.TimedRotatingFileHandler(logger_path, 'midnight', 1, 0)
        file_handler.setFormatter(Formatter(format))
        file_handler.suffix = "%Y%m%d-%H%M.log"
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(Formatter(format))
        logger.addHandler(file_handler)
        logger.addHandler(console_handler)
        logger.setLevel(logging.INFO)


    def get_video_size(self):
        for i in range(3):
            try:
                response = requests.get(self.url, stream=True)
                content_length = int(response.headers.get("content-length", None))
                if content_length:
                    self.video_size = content_length
                    return content_length
                else:
                    logger.info(f"content-length={content_length}")
            except Exception as e:
                continue
        return None

    def split_video_block(self, video_size):
        """
            根据content-length分割视频文件
        :return:
        """
        if video_size <= 0:
            return None

        split_num = (video_size) // self.block_size + 1
        # 分割文件，将分割结果加入下载队列中
        sections = {}
        for i in range(split_num):
            if i + 1 != split_num:
                sections[str(i)] = (i * self.block_size, (i + 1) * self.block_size - 1)
            else:
                sections[str(i)] = (i * self.block_size, video_size -1 )
        return sections

    def add_video_block_to_queue(self, section_dict):
        self.section_count = len(section_dict.keys())
        for id in list(section_dict.keys()):
            section = section_dict.pop(id)
            # 将文件片段加入到待下载队列中
            self.q_file_section.put({"id": id, "size": int(section[1] - section[0] + 1), "section": f"bytes={section[0]}-{section[1]}"})

    def download_executor(self):
        with ThreadPoolExecutor(max_workers=self.max_threads) as executor:
            for i in range(self.section_count):
                section = self.q_file_section.get()
                executor.submit(self.download, section)

    def download(self, section_sub_dict):
        section_size = section_sub_dict.get("size")
        section = section_sub_dict.get("section")
        split_id = section_sub_dict.get("id")

        sub_file_path = f"{self.output_path}/{split_id}.mp4"
        if os.path.exists(sub_file_path):
            logger.info(f"{sub_file_path} exists")
            return

        # 构建请求头
        headers = {"Range": section}
        logger.info(f"block_{split_id} start download path={sub_file_path}")
        stream_content = requests.get(url=self.url, headers=headers, stream=True).content
        stream_size = len(stream_content)
        logger.info(f"block_{split_id} end download stream_size={stream_size} section_size={section_size}")

        if section_size == stream_size:
            with open(sub_file_path, mode="wb") as f:
                f.write(stream_content)
            logger.info(f"block_{split_id} write complete")

    def check_section(self):
        file_size = 0
        # 校验分割文件数是否合理
        if self.section_count == 0:
            return False
        # 校验分割文件是否存在
        for i in range(self.section_count):
            sub_file_path = f"{self.output_path}/{i}.mp4"
            if os.path.exists(sub_file_path):
                file_size += os.path.getsize(sub_file_path)
            else:
                return False

        return file_size == self.video_size


    def check_and_merge_section(self):
        check_status = self.check_section()
        logger.info(f"{self.file_path} check_block_status={check_status}")
        if check_status:
            if self.section_count == 1:
                os.rename(f"{self.output_path}/0.mp4", self.file_path)
                return True
            else:
                # 文件合并
                with open(self.file_path, "ab") as result_file:
                    for i in range(self.section_count):
                        sub_file_path = f"{self.output_path}/{i}.mp4"
                        with open(sub_file_path, "rb") as f:
                            result_file.write(f.read())
                return True
        else:
            return False

    def clear_video_block(self):
        if self.remove_block_file and self.section_count > 1:
            for i in range(self.section_count):
                sub_file_path = f"{self.output_path}/{i}.mp4"
                os.remove(sub_file_path)

    def run(self):
        try:
            logger.info(f"{self.file_path} start download ")
            if not os.path.exists(self.output_path):
                os.makedirs(self.output_path)

            # 获取视频总长度
            video_size = self.get_video_size()
            if video_size and os.path.exists(self.file_path) and os.path.getsize(self.file_path) == video_size:
                logger.info(f"{self.file_path} exists")
                return self.file_path

            # 对视频长度进行分割
            section_dict = self.split_video_block(video_size)
            logger.info(f"{self.file_path} split_video_block len={len(section_dict)}")
            # 把分割结果推入任务队列
            self.add_video_block_to_queue(section_dict)
            # 开始多线程分块下载
            self.download_executor()

            # 视频合并
            flag = self.check_and_merge_section()

            # 清理中间文件
            if not flag:
                return None
            else:
                self.clear_video_block()
        except Exception as e:
            os.remove(self.output_path)
            return None
        return self.file_path


if __name__ == '__main__':
    start_time = int(time.time() * 1000)

    url = "$视频可下载链接"
    video_key = "sirius/video_key.mp4"
    down_loader = BlockDownLoader(url, video_key, block_size=1024 * 1024 * 5)
    result_path = down_loader.run()

    end_time = int(time.time() * 1000)
    cost_time = end_time - start_time
    logging.info(f"result = {result_path}, cost time: {str(cost_time)} ms")