Python 多线程文件下载器 - 高效下载文件 - 常规

import sys
import os
import math
import time
import requests
import urllib.request
import threading
from tqdm import tqdm


class DownloadThread(threading.Thread):
    def __init__(self, url, start, end, fileobj, bar):
        threading.Thread.__init__(self)
        self.url = url
        self.start = start
        self.end = end
        self.fileobj = fileobj
        self.bar = bar

    def run(self):
        headers = {'Range': 'bytes=%d-%d' % (self.start, self.end)}
        response = requests.get(self.url, headers=headers, stream=True)
        for chunk in response.iter_content(chunk_size=1024):
            if not chunk:
                break
            self.fileobj.write(chunk)
            self.bar.update(len(chunk))


def download_file(url:str, filename:str, num_threads:int=4) -> None:
    response = requests.get(url, stream=True)
    if 'Content-Length' in response.headers:
        file_size = int(response.headers['Content-Length'])
    else:
        print('The file size is unknown, use single thread download.')
        download_single_thread(url=url, filename=filename)
        return

    # 判断文件大小是否为0
    if file_size == 0:
        raise ValueError('Failed to get file size, please check the download URL')
    
    with open(filename, 'wb') as file:
        with tqdm(total=file_size, unit='B', unit_scale=True,
                  desc=os.path.basename(filename), ncols=80, position=0) as bar:

            download_time = 0
            chunk_size = 1024
            avg_speed = 0
            last_speed = 0
            while True:
                start = time.time()
                chunk = response.raw.read(chunk_size*num_threads)
                end = time.time()

                if not chunk:
                    break

                used_time = end - start
                download_time += used_time

                speed = chunk_size*num_threads/used_time
                avg_speed = (avg_speed + speed) / 2

                if avg_speed > 0:
                    last_speed = avg_speed
                    num_threads = int(file_size / avg_speed / chunk_size) + 1
                num_threads = min(max(num_threads, 1), 32)

                thread_pool = []
                start = 0
                end = -1
                for i in range(num_threads):
                    start = end + 1
                    end = start + chunk_size - 1

                    if i == num_threads - 1:
                        end = file_size - 1

                    thread = DownloadThread(url=url, start=start, end=end, fileobj=file, bar=bar)
                    thread.setDaemon(True)
                    thread_pool.append(thread)

                for thread in thread_pool:
                    thread.start()

                for thread in thread_pool:
                    thread.join()

                bar.update(chunk_size*num_threads)

            bar.close()

    return filename


def download_single_thread(url:str, filename:str) -> None:
    urllib.request.urlretrieve(url, filename)
    return filename


def download(input_str:str, filename:str, num_threads:int=4) -> None:
    try:
        if os.path.isdir(filename):
            output_filename = os.path.join(filename, os.path.basename(input_str))
        else:
            output_filename = filename
        
        # if input_str.startswith('magnet:'):
        #     download_magnet(input_str, output_filename)
        # else:
        download_file(input_str, output_filename, num_threads)
        
        print(f'Download succeed! The file is saved as: {output_filename}')
    except Exception as e:
        print('Download failed, exception:', e)


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print('Usage: python download.py url output_file')
        sys.exit(1)

    url = sys.argv[1]
    filename = sys.argv[2]

    download(input_str=url, filename=filename, num_threads=4)

代码说明:

DownloadThread 类:
- 继承 threading.Thread 类，用于创建下载线程。
- 初始化时接收 url、起始字节位置、结束字节位置、文件对象和进度条对象。
- run 方法负责从服务器下载对应字节范围的数据并写入文件对象，并更新进度条。
download_file 函数:
- 获取文件大小，并使用 tqdm 创建进度条。
- 使用循环不断读取数据并创建下载线程，并将线程添加到线程池中。
- 启动所有线程并等待所有线程执行完毕。
- 动态调整线程数量以优化下载速度。
download_single_thread 函数:
- 使用 urllib.request.urlretrieve 函数进行单线程下载。
download 函数:
- 接收下载链接和保存路径，并根据链接类型选择不同的下载方式。
- 如果链接为磁力链接，则调用 download_magnet 函数进行下载 (需要自行实现)。
- 如果链接为普通链接，则调用 download_file 函数进行下载。
主函数:
- 从命令行参数获取下载链接和保存路径，并调用 download 函数进行下载。

使用方法:

将代码保存为 Python 文件，例如 download.py。
打开命令行，进入代码所在目录。
执行命令 python download.py <下载链接> <保存路径>，例如 python download.py https://www.example.com/file.zip /home/user/downloads。

注意:

代码中没有实现下载磁力链接的功能，需要自行实现。
下载速度受网络带宽、服务器性能等因素影响。
可以根据实际情况调整线程数量和 chunk_size 的值以优化下载速度。
建议在使用该代码之前阅读相关文档了解其原理和使用方法。