Butterchalk/main.py

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

""" 
---------------------------------------
 # @Project    : Alcorart Downloader
 # @File       : main.py
 # @Author     : GrayZhao
 # @Date       : 2023/2/18 17:04
 # @Version    : 
 # @Description : 
---------------------------------------
"""
from gevent import monkey; monkey.patch_all()
import gevent
import requests
import os
from queue import Queue
from tqdm import tqdm

URL = "https://baraag.net/api/v1/accounts/387484/statuses?only_media=true&limit=40"
HEADER = {
    "referer": "https://baraag.net/@Butterchalk/media",
    "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36 Edg/110.0.1587.46"
}


def producer(until_id: str, data_q: Queue):
    """
    下载进程

    :param until_id: 截止作品ID（包含该ID）
    :param data_q: 下载队列
    """
    flag = True
    img_index = 0
    work_url = URL
    while flag:
        now_id = ""
        response = requests.get(url=work_url, headers=HEADER)
        data_json = response.json()
        for data_item in data_json:
            img_index += 1
            image_url = data_item["media_attachments"][0]["url"]
            data_q.put_nowait({"img_index": img_index, "image_url": image_url})
            now_id = data_item["id"]
            if now_id == until_id:
                flag = False
                break
        work_url = URL + f"&max_id={now_id}"


def download(data_q: Queue):
    while not data_q.empty():
        image_info = data_q.get_nowait()
        image_index = image_info["img_index"]
        image_url = image_info["image_url"]
        ext = os.path.splitext(image_url)[1]
        response = requests.get(url=image_url, headers=HEADER, stream=True)
        total_size = int(response.headers.get("content-length", 0))
        block_size = 1024
        save_path = f"img/{image_index}{ext}"

        with tqdm(desc=f"第{image_index}张图片", total=total_size, unit="KB",
             unit_scale=True, leave=False) as pbar:
            with open(save_path, "wb") as file:
                for data in response.iter_content(block_size):
                    file.write(data)
                    pbar.update(len(data))


def consumer(data_q: Queue):
    tasks = [gevent.spawn(download, data_q) for x in range(4)]
    gevent.joinall(tasks)


if __name__ == '__main__':
    data_queue = Queue()
    ID = "108186424765846031"
    producer(ID, data_queue)
    # 并发下载
    print(f"共有{data_queue.qsize()}张图片，开始下载... [请耐心等待]")
    consumer(data_queue)
    print("下载完成...")