Python之requests模塊-大文件分片上傳

最近在做接口測試時，拿到一個分片上傳文件的接口，http接口請求頭中的Content-Type為multipart/form-data。需要在客戶端將大文件分片成數據塊后，依次傳給服務端，由服務端還原成大文件，此外，為了確保傳輸后的數據是完整的，客戶端會在分片前，根據原文件生成md5值並被攜帶在每次的http請求中，服務端在還原文件后會進行校驗。

如何使用requests模塊，實現上述接口測試的需求呢？首先，需要將問題分解：

1. requests如何傳輸Content-Type為multipart/form-data的數據？
2. 如何根據原文件生成md5值？
3. 如何將大文件分片成數據塊？

本文將逐一為大家解答。

發送multipart/form-data請求

這里需要用到輔助庫requests_toolbelt，使用MultipartEncoder類創建一個multipart/form-data類型的data充當請求體。此外，我們請求頭中的Content-Type除了multipart/form-data還需生成boundary，如下例所示：

import requests
from requests_toolbelt import MultipartEncoder
import os


def upload_multipart(url, file_path):
    filename = file_path.split("\\")[-1:][0]
    total_size = os.path.getsize(file_path)
    data = MultipartEncoder(
        fields={
            "filename": filename,
            "totalSize": str(total_size),
            "file": (filename, open(file_path, 'rb'), 'application/octet-stream')
        }
    )
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.107 Safari/537.36",
        "Accept": "application/json",
        "Accept-Encoding": "gzip, deflate",
        "Connection": "keep-alive",
        "Content-Type": data.content_type
    }
    with requests.post(url, headers=headers, data=data) as response:
        assert response.status_code == 200

根據原文件生成md5值

使用hashlib庫，如下例所示：

import hashlib


def get_md5(path):
    m = hashlib.md5()
    with open(path, 'rb') as f:
        for line in f:
            m.update(line)
    md5code = m.hexdigest()
    return md5code

大文件分片成數據塊

如下例所示，定義數據塊的大小為2MB，根據文件大小划分出數據塊的總數量，通過fileObject.seek()函數偏移文件的指針到當前數據塊的位置，依次讀取數據塊並發送請求，每個請求都帶上了md5值。

import requests
from requests_toolbelt import MultipartEncoder
import os
import math


def upload_slice_file(url, file_path):
    chunk_size = 1024*1024*2
    filename = file_path.split("\\")[-1:][0]
    total_size = os.path.getsize(file_path)
    current_chunk = 1
    total_chunk = math.ceil(total_size/chunk_size)

    while current_chunk <= total_chunk:
        start = (current_chunk - 1)*chunk_size
        end = min(total_size, start+chunk_size)
        with open(file_path, 'rb') as f:
            f.seek(start)
            file_chunk_data = f.read(end-start)
        data = MultipartEncoder(
            fields={
                "filename": filename,
                "totalSize": str(total_size),
                "currentChunk": str(current_chunk),
                "totalChunk": str(total_chunk),
                "md5": get_md5(file_path),
                "file": (filename, file_chunk_data, 'application/octet-stream')
            }
        )
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.107 Safari/537.36",
            "Accept": "application/json",
            "Accept-Encoding": "gzip, deflate",
            "Connection": "keep-alive",
            "Content-Type": data.content_type
        }
        with requests.post(url, headers=headers, data=data) as response:
            assert response.status_code == 200

        current_chunk = current_chunk + 1