【Python】圧縮されたS3オブジェクトをローカルに保存せず読み込む

S3に保存されたgzip/zip圧縮ファイルを、ローカルファイルとして保存することなく読み込みたい。

前提条件

• Python: 3.7.4
• AWS SDK: boto3-1.9.230以上

ポイント

• 標準のgzip, zipfileパッケージを使ってファイルを読み込んでファイルオブジェクトに変換する。
  

（bzip2とかは考えなくて良いんじゃないかな）

つまり、以下のようにして読み込んだ際と同様に扱いたい。

import gzip
gz_file = 'path/to/file.csv.gz'
file = gzip.open(gz_file, 'rt')

file.read()
..

パッケージ毎にファイルオブジェクトの変換方法に癖があるので、それぞれ対応する。
ioパッケージのBytesIOやらTextIOWrapperを使う。
ついでに、圧縮形式も正しいかチェックする。

コード

gzip圧縮の場合

import boto3
import gzip
import io


# 圧縮形式が正しいかチェック
def valid_gzip_format(obj):
    file = gzip.open(io.BytesIO(obj), 'rt')
    try:
        file.read()
        return True
    except OSError:
        return False


s3 = boto3.client('s3')
bucket = 'bucket_name'
key = 'path/to/file.csv.gz'
# S3オブジェクトの取得
obj = s3.get_object(
    Bucket=bucket,
    Key=key
)['Body'].read()

# main
if valid_gzip_format(obj):
    file = gzip.open(io.BytesIO(obj), 'rt')
else:
    raise InvalidCompressError('ファイル拡張子(.gz)に対して圧縮形式が正しくありません。')

# 中身確認
for row in file.readlines():
    print(row.replace('\n', ''))

zip圧縮の場合

import boto3
import io
import zipfile


# 圧縮形式が正しいかチェック
def valid_zip_format(obj):
    return zipfile.is_zipfile(io.BytesIO(obj))


s3 = boto3.client('s3')
bucket = 'bucket_name'
key = 'path/to/file.csv.zip'
# S3オブジェクトの取得
obj = s3.get_object(
    Bucket=bucket,
    Key=key
)['Body'].read()

# main
if valid_zip_format(obj):
    zf = zipfile.ZipFile(io.BytesIO(obj))
    if len(zf.namelist()) > 1:
        raise InvalidCompressError('zipアーカイブの中にファイルが複数あります。')
    file = io.TextIOWrapper(zf.open(zf.namelist()[0], 'r'))
else:
    raise InvalidCompressError('ファイル拡張子(.zip)に対して圧縮形式が正しくありません。')

# 中身確認
for row in file.readlines():
    print(row.replace('\n', ''))

zipfileではopenのモードに'rt'が無いのでTextIOWrapperを使う必要がある。
参考 - ZIP アーカイブの処理 — Python 3.7.5rc1 ドキュメント

バージョン 3.6 で変更: Removed support of mode='U'. Use io.TextIOWrapper for reading compressed text files in universal newlines mode.

まとめ

上記のコードで取得したS3オブジェクトはbytes型で保持されるだけなので、with構文で扱う必要はないと思う（合ってる？）。
あと実際に中身を読み込む際に、末尾に改行コードが含まれるのでそのあたりも考慮が必要（もろちん、バイナリモードでも含まれるけど）。