使用 Python 從 Google Drive 下載檔案

import os
import re
import requests
from tqdm import tqdm
from bs4 import BeautifulSoup

def download_from_google(file_id: str, file_name: str, target: str = "."):
    """
    Downloads a file from Google Drive, handling potential confirmation tokens for large files.

    Args:
        file_id (str):
            The ID of the file to download from Google Drive.
        file_name (str):
            The name to save the downloaded file as.
        target (str, optional):
            The directory to save the file in. Defaults to the current directory (".").

    Raises:
        Exception: If the download fails or the file cannot be created.

    Notes:
        This function handles both small and large files. For large files, it automatically processes
        Google's confirmation token to bypass warnings about virus scans or file size limits.

    Example:
        Download a file to the current directory:
            download_from_google(
                file_id="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
                file_name="example_file.txt"
            )

        Download a file to a specific directory:
            download_from_google(
                file_id="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
                file_name="example_file.txt",
                target="./downloads"
            )
    """
    # 第一次嘗試：docs.google.com/uc?export=download&id=檔案ID
    base_url = "https://docs.google.com/uc"
    session = requests.Session()
    params = {
        "export": "download",
        "id": file_id
    }
    response = session.get(base_url, params=params, stream=True)

    # 如果已經出現 Content-Disposition，代表直接拿到檔案
    if "content-disposition" not in response.headers:
        # 先嘗試從 cookies 拿 token
        token = None
        for k, v in response.cookies.items():
            if k.startswith("download_warning"):
                token = v
                break

        # 如果 cookies 沒有，就從 HTML 解析
        if not token:
            soup = BeautifulSoup(response.text, "html.parser")
            # 常見情況：HTML 裡面有一個 form#download-form
            download_form = soup.find("form", {"id": "download-form"})
            if download_form and download_form.get("action"):
                # 將 action 裡的網址抓出來，可能是 drive.usercontent.google.com/download
                download_url = download_form["action"]
                # 收集所有 hidden 欄位
                hidden_inputs = download_form.find_all(
                    "input", {"type": "hidden"})
                form_params = {}
                for inp in hidden_inputs:
                    if inp.get("name") and inp.get("value") is not None:
                        form_params[inp["name"]] = inp["value"]

                # 用這些參數去重新 GET
                # 注意：原本 action 可能只是相對路徑，這裡直接用完整網址
                response = session.get(
                    download_url, params=form_params, stream=True)
            else:
                # 或者有些情況是直接在 HTML 裡 search confirm=xxx
                match = re.search(r'confirm=([0-9A-Za-z-_]+)', response.text)
                if match:
                    token = match.group(1)
                    # 帶上 confirm token 再重新請求 docs.google.com
                    params["confirm"] = token
                    response = session.get(
                        base_url, params=params, stream=True)
                else:
                    raise Exception("無法在回應中找到下載連結或確認參數，下載失敗。")

        else:
            # 直接帶上 cookies 抓到的 token 再打一次
            params["confirm"] = token
            response = session.get(base_url, params=params, stream=True)

    # 確保下載目錄存在
    os.makedirs(target, exist_ok=True)
    file_path = os.path.join(target, file_name)

    # 開始把檔案 chunk 寫到本地，附帶進度條
    try:
        total_size = int(response.headers.get('content-length', 0))
        with open(file_path, "wb") as f, tqdm(
            desc=file_name,
            total=total_size,
            unit="B",
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for chunk in response.iter_content(chunk_size=32768):
                if chunk:
                    f.write(chunk)
                    bar.update(len(chunk))

        print(f"File successfully downloaded to: {file_path}")

    except Exception as e:
        raise Exception(f"File download failed: {e}")

download_from_google(
    file_id="YOUR_FILE_ID",
    file_name="big_model.onnx",
    target="./models"
)

from download_from_google import download_from_google
import argparse

def main():
    parser = argparse.ArgumentParser(description="Download files from Google Drive.")
    parser.add_argument("--file-id", required=True, help="Google Drive file ID.")
    parser.add_argument("--file-name", required=True, help="Output file name.")
    parser.add_argument("--target", default=".", help="Output directory. Defaults to current folder.")
    args = parser.parse_args()

    download_from_google(file_id=args.file_id, file_name=args.file_name, target=args.target)

if __name__ == "__main__":
    main()