中野智文のブログ

データ・マエショリストのメモ

s3上のgzip圧縮されたファイルの中身をイテレータで取得する

背景

s3上のgzip圧縮されたファイルの中身をイテレータで取得する方法がなかなか見つからない。

コード

import boto3
import tempfile
import gzip
import datetime as dt
from typing import Iterator


def s3_gz_cat(bucket: str, prefix='') -> Iterator[bytes]:
    s3 = boto3.resource('s3')
    s3_bucket = s3.Bucket(bucket)
    for s3_obj_sum in s3_bucket.objects.filter(Prefix=prefix):
        # s3_obj_sum.key がファイル名(path)なので、末尾が .gz 判定とかしたい場合はここで。                                                                                                                                                                   
        with tempfile.TemporaryFile() as temp:
            s3_obj = s3_obj_sum.Object()
            s3_obj.download_fileobj(temp)
            temp.seek(0)
            with gzip.GzipFile(fileobj=temp) as gz:
                for b in gz:
                    yield b


# 以下は使い方

if __name__ == '__main__':
    # my_bucket 内全て
    for b in s3_gz_cat(bucket='my_bucket'):
        print(b)

    # my_dir 内に限定
    for b in s3_gz_cat(bucket='my_bucket', prefix='my_dir'):
        print(b)

IOバッファでやる方法もあるみたいだが、バッファが溢れたときどうなるのか(詰まったりする?)心配なので、tempfile でやったほうがいいような気がする。