Pythonの io.BytesIO と zipfile.ZipFile の組み合わせ：コーナーケース、またはZIPファイルを書き込む場合

with文は1行で書ける。しかし…。

以前のエントリで、Pythonの io.BytesIO と zipfile.ZipFile の組み合わせについて扱った。

xaro.hatenablog.jp

その際の結論として、io.BytesIO を明示的に閉じるように勧めた。

with io.BytesIO(content) as bs:
    with zipfile.ZipFile(bs) as zf:
        names: list[str] = zf.namelist()

さて、 Pythonの with 文は複数のコンテキストマネージャを扱える。 よって、上記のスニペットは次のようにも書ける。

with io.BytesIO(content) as bs, zipfile.ZipFile(bs) as zf:
    names: list[str] = zf.namelist()

with 文がある行は多少長くなるものの、インデントレベルを削減できるメリットもある。

しかし、1行にまとめる際に気を付けるポイントがある。 ZIPファイルの書き込みを行う場合だ。

やや人工的な例だが、ZIPファイルの書き込みのコード例を以下に記載する。 イメージとしては、ストリーム上にZIPファイルを作成して、Djangoの SimpleUploadedFile オブジェクトを作成する、Boto3でストリーム経由でS3にアップロードする、といった状況を模したものである。

import io
import zipfile


def make_zip_stream() -> bytes:
    with io.BytesIO() as bytes_stream, zipfile.ZipFile(bytes_stream, mode="w") as zip_stream:
        zip_stream.write("sample.txt")
        return bytes_stream.getvalue()


with zipfile.ZipFile(io.BytesIO(make_zip_stream()), "r") as zf:
    print(zf.testzip())

さて、 make_zip_stream() 関数は io.BytesIO の 中身を丸ごと返す形をとっている。 この際に、 zipfile.ZipFile は閉じられるのだろうか。

zipfile.ZipFileの公式ドキュメントのclose()の節には次のように書かれている。

アーカイブファイルをクローズします。close() はプログラムを終了する前に必ず呼び出さなければなりません。さもないとアーカイブ上の重要なレコードが書き込まれません。

つまり、return bytes_stream.getvalue() をする前にzipfile.ZipFileを閉じないといけない。 果たして、上記のコードはどうなるだろうか。

Traceback (most recent call last):
  File "zip_and_bytesio.py", line 12, in <module>
    with zipfile.ZipFile(io.BytesIO(make_zip_stream()), "r") as zf:
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    raise BadZipFile("File is not a zip file")
zipfile.BadZipFile: File is not a zip file

zipfile.BadZipFile 例外が発生した。 処理としては、with文のドキュメントにある通り、with文の中身（ここでは zip_stream.write("sample.txt") と return bytes_stream.getvalue()）が実行されて、その後にコンテキストマネージャの __exit__() が実行される。 つまり、ZIPファイルがクローズされる前にバイナリストリームの中身を取得するため、「アーカイブ上の重要なレコードが書き込まれ」ない状態に陥る。

どうすればいいか。 単にクローズすればよいのである。

import io
import zipfile


def make_zip_stream() -> bytes:
    with io.BytesIO() as bytes_stream, zipfile.ZipFile(bytes_stream, mode="w") as zip_stream:
        zip_stream.write("sample.txt")
        zip_stream.close()  # ここで明示的にcloseする
        return bytes_stream.getvalue()


with zipfile.ZipFile(io.BytesIO(make_zip_stream()), "r") as zf:
    print(zf.testzip())

クローズしたファイルをクローズする分には問題はない。 zf.testzip() は、成功した場合は None を返す点に注意しよう。

with 文を使っておけばリソース問題は解決、とは必ずしもならない。 やや人工的なコーナーケースだが、io.BytesIOを駆使したコードだと案外遭遇しやすい例なのかもしれない。