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 = "."):
    """
    Google Driveからファイルをダウンロードし、大きなファイルのために確認トークンを処理します。

    引数：
        file_id (str):
            Google DriveからダウンロードするファイルのID。
        file_name (str):
            ダウンロードしたファイルの保存名。
        target (str, optional):
            ファイルを保存するディレクトリ。デフォルトはカレントディレクトリ（"."）。

    例外：
        Exception: ダウンロードに失敗した場合や、ファイルを作成できない場合。

    備考：
        この関数は小さなファイルと大きなファイルの両方を処理します。大きなファイルの場合、Googleの確認トークンを自動的に処理し、ウイルススキャン警告やファイルサイズ制限を回避します。

    例：
        カレントディレクトリにファイルをダウンロード：
            download_from_google(
                file_id="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
                file_name="example_file.txt"
            )

        特定のディレクトリにファイルをダウンロード：
            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:
        # 最初にクッキーからトークンを取得
        token = None
        for k, v in response.cookies.items():
            if k.startswith("download_warning"):
                token = v
                break

        # クッキーにトークンがない場合は、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内のURLを取得、通常は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は相対パスの可能性があるので、ここでは完全なURLを使用
                response = session.get(download_url, params=form_params, stream=True)
            else:
                # もしくはHTML内に直接confirm=xxxが含まれていることもある
                match = re.search(r'confirm=([0-9A-Za-z-_]+)', response.text)
                if match:
                    token = match.group(1)
                    # confirmトークンを含めて再度docs.google.comにリクエスト
                    params["confirm"] = token
                    response = session.get(base_url, params=params, stream=True)
                else:
                    raise Exception("レスポンス内でダウンロードリンクまたは確認パラメータが見つかりませんでした。ダウンロード失敗。")

        else:
            # クッキーから取得したトークンを使って再リクエスト
            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_path}")

    except Exception as e:
        raise Exception(f"ファイルのダウンロードに失敗しました: {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="Google Driveからファイルをダウンロードします。")
    parser.add_argument("--file-id", required=True, help="Google DriveのファイルID。")
    parser.add_argument("--file-name", required=True, help="出力ファイル名。")
    parser.add_argument("--target", default=".", help="出力ディレクトリ。デフォルトはカレントディレクトリ。")
    args = parser.parse_args()

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

if __name__ == "__main__":
    main()