🎃

InquirerPyを使って、S3から対話式でファイルをダウンロードする

2023/07/20に公開

Python

inquirerpy

tech

はじめに

今回はpythonのInquirerPyを使って、
S3から対話式でわかりやすくファイルをダウンロード出来るCLIツールを作りました。

よかったら見てみて下さい。

動作イメージ

とても便利です！

コード

main.py

import boto3
from InquirerPy import inquirer

s3 = boto3.resource('s3')

def list_buckets():
    return [bucket.name for bucket in s3.buckets.all()]

def list_objects(bucket_name, prefix=''):
    bucket = s3.Bucket(bucket_name)
    all_objects = [obj.key for obj in bucket.objects.filter(Prefix=prefix)]
    if prefix:
        all_objects = [obj.replace(prefix, '', 1) for obj in all_objects]
    return [obj for obj in all_objects if obj != ''] 

def download_file(bucket_name, file_key, download_path):
    s3.meta.client.download_file(bucket_name, file_key, download_path)
    print(f'File downloaded at {download_path}')

def select_and_download():
    selected_bucket = inquirer.select(
        message="Select an S3 bucket:",
        choices=list_buckets(),
    ).execute()
    
    current_prefix = ''
    while True:
        objects = list_objects(selected_bucket, current_prefix)
        if not objects:
            print("No files found.")
            break

        selected_obj = current_prefix + inquirer.select(
            message="Select a file or a directory:",
            choices=objects,
        ).execute()

        if selected_obj.endswith('/'):  # Directory
            current_prefix = selected_obj
        else:  # File
            download_path = input(f'Enter the download path for {selected_obj}: ')
            download_file(selected_bucket, selected_obj, download_path)
            break

if __name__ == "__main__":
    select_and_download()

流れ

s3バケットを選択する
s3バケットのオブジェクトを選択する。ファイルを選べばダウンロードへ、フォルダを選択したらそのフォルダを起点にして再度オブジェクトを選択する
ダウンロードする際のファイル名を指定する
ダウンロードできたことを確認する

環境

python 3.10.9
boto3 1.28.7
botocore 1.31.7
inquirerpy 0.3.4

ポイント

現在のディレクトリを選択肢から除外する

def list_objects(bucket_name, prefix=''):
    bucket = s3.Bucket(bucket_name)
    all_objects = [obj.key for obj in bucket.objects.filter(Prefix=prefix)]
    if prefix:
        all_objects = [obj.replace(prefix, '', 1) for obj in all_objects]
    return [obj for obj in all_objects if obj != '']

指定したs3バケット内のオブジェクトを一覧出力する関数です。
if prefixとしているのは、各オブジェクトを一覧表示した際に、ディレクトリ名を除外するためです。
このようにしないと、prefixで参照しているディレクトリ名が延々表示され、無限ループのようになります。

フォルダ1を選択
フォルダ1を選択できてしまう。フォルダ1を選択
以下繰り返し

質問を作る

def select_and_download():
    selected_bucket = inquirer.select(
        message="Select an S3 bucket:",
        choices=list_buckets(),
    ).execute()

inquirer.selectのchoiceで選択肢として表示する内容を指定し、executeで質問を投げかけます。
今回はlist_buckets()関数を使用してs3バケットを一覧表示し、そこから選択させるように促しています。

ファイルをダウンロードするまでファイル一覧出力を続ける

    while True:
        objects = list_objects(selected_bucket, current_prefix)
        if not objects:
            print("No files found.")
            break

        selected_obj = current_prefix + inquirer.select(
            message="Select a file or a directory:",
            choices=objects,
        ).execute()

        if selected_obj.endswith('/'):  # Directory
            current_prefix = selected_obj
        else:  # File
            download_path = input(f'Enter the download path for {selected_obj}: ')
            download_file(selected_bucket, selected_obj, download_path)
            break

while Trueとすることで、breakが出現するまでコードをループするようにしています。
breakはファイルが見つからない場合、ファイルをダウンロードした場合に現れますので、それまではディレクトリの中身を一覧表示する動作を続けます。
結果的にディレクトリを深ぼっていく形となります。

おわりに

pythonのInquirerPy、わかりやすくてこんな簡単にCLIツールが作成できます。(まあChatGPTもいるのですが...)
みなさまもよかったら使ってみて下さい。