【Python】指定のS3ディレクトリから再帰的にファイルをダウンロード

自分用のメモです🦔

Amazon S3のファイルを再帰的にダウンロードしてローカルに保存するスクリプトです。

アクセスキーとかも雑にスクリプト上で指定しています。

s3_download.py

import os
import argparse

from typing import List, Any

import boto3  # type: ignore

# AWS認証情報の設定
os.environ['AWS_DEFAULT_REGION'] = 'ap-northeast-1'
os.environ['AWS_ACCESS_KEY_ID'] = ''
os.environ['AWS_SECRET_ACCESS_KEY'] = ''


def get_args() -> Any:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--bucket_name',
        type=str,
        default='',
        help='対象のS3バケット名を指定',
    )
    parser.add_argument(
        '--directory',
        type=str,
        default='',
        help='対象のディレクトリを指定',
    )

    args = parser.parse_args()

    return args


def get_path_list(
    bucket: str,
    prefix: str,
    recursive: bool = False,
) -> List[str]:
    # 指定されたS3バケットとディレクトリから、パスのリストを取得
    path_list: List[str] = []
    path_list = _get_path_list(bucket, prefix, recursive=recursive)

    return path_list


def _get_path_list(
    bucket: str,
    prefix: str,
    keys: Any = None,
    marker: str = '',
    recursive: bool = False,
) -> List[str]:
    s3 = boto3.client('s3')

    # 再帰的にファイルリストを取得するかどうかで、API呼び出しを変更
    if recursive:
        response = s3.list_objects(
            Bucket=bucket,
            Prefix=prefix,
            Marker=marker,
        )
    else:
        response = s3.list_objects(
            Bucket=bucket,
            Prefix=prefix,
            Marker=marker,
            Delimiter='/',
        )

    if keys is None:
        keys = []

    # S3レスポンスからファイルやディレクトリのパスを取得
    if 'CommonPrefixes' in response:
        keys.extend(
            [content['Prefix'] for content in response['CommonPrefixes']])
    elif 'Contents' in response:
        keys.extend([content['Key'] for content in response['Contents']])
        if 'IsTruncated' in response:
            return _get_path_list(
                bucket=bucket,
                prefix=prefix,
                keys=keys,
                marker=keys[-1],
                recursive=recursive,
            )

    return keys


def download(local_file_name: str, s3_bucket: str, s3_object_key: str):
    print(s3_object_key)

    s3_client = boto3.client('s3')
    with open(local_file_name, 'wb') as f:
        s3_client.download_fileobj(
            s3_bucket,
            s3_object_key,
            f,
            Callback=None,
        )


def main() -> None:
    args = get_args()

    s3_bucket_name = args.bucket_name
    s3_directory = args.directory

    # ファイルリスト取得
    file_list = get_path_list(
        s3_bucket_name,
        s3_directory,
        recursive=True,
    )

    for file_path in file_list:
        basename = os.path.basename(file_path)
        directory_path = os.path.dirname(file_path)
        if basename == '':
            continue

        # バケット名とディレクトリに基づいて、ローカルのディレクトリを作成
        os.makedirs(s3_bucket_name + '/' + directory_path, exist_ok=True)

        # ダウンロード
        download(
            s3_bucket_name + '/' + file_path,
            s3_bucket_name,
            file_path,
        )


if __name__ == '__main__':
    main()