Pythonのio.BytesIOとzipfile.ZipFileの組み合わせ¶

2024-03-22 公開

インターネット経由で取得したZIPファイルを手元で加工する、という状況を考える。たとえば、次のようなコードを書いたとする。 URL は ZIPファイルを取得できるものならば何でもよいが、今回は環境に配慮してローカルに1回だけダウンロードしてhttp.serverで簡易Webサーバを立てることでキャッシュしている。

import io
import urllib.request
import zipfile

URL = "http://localhost:8000/python-3.12.2-embed-amd64.zip"

with urllib.request.urlopen(URL) as f:
    content: bytes = f.read()

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

for name in names:
    print(name)

urllib.request.urlopen() 経由で取得したものはbytesであり、ファイルオブジェクトではない。 zipfile.ZipFile の第一引数file はファイルのパス、ファイルオブジェクト、 pathlib.Path のいずれかである必要があるため、urllib.request.urlopen() 経由で取得したものをそのまま渡すわけにはいかない。そのため、 io.Bytes() を経由することでインメモリーストリームとして渡す。

このエントリの主題は、上記のコードにおける io.BytesIO(content) の取り扱いについてである。 with zipfile.ZipFile(io.BytesIO(content)) as zf: と記述しているので、 names: list[str] = zf.namelist() が完了したら zipfile.ZipFile(io.BytesIO(content)) の close() メソッドが呼ばれる。その際に、中身である io.BytesIO(content) は閉じられるのだろうか。

公式ドキュメントとソースコードを巡る冒険¶

公式ドキュメントには

The buffer is discarded when the close() method is called.

とある通り、 BytesIO の close() メソッドが呼ばれなければインメモリーストリームのバッファは解放されない。

まず、ソースコードにある __init__() から説明に必要な個所を抜き出す。

class ZipFile:
    fp = None                   # Set here since __del__ checks it

    def __init__(self, file, mode="r", compression=ZIP_STORED, allowZip64=True,
                 compresslevel=None, *, strict_timestamps=True, metadata_encoding=None):
        """Open the ZIP file with mode read 'r', write 'w', exclusive create 'x',
        or append 'a'."""
        ...

        # Check if we were passed a file-like object
        if isinstance(file, os.PathLike):
            file = os.fspath(file)
        if isinstance(file, str):
            # No, it's a filename
           ...
        else:
            self._filePassed = 1
            self.fp = file
            self.filename = getattr(file, 'name', None)
        self._fileRefCnt = 1
        ...

ここで認識してほしいのは以下の3点である。

zipfile.ZipFile の第一引数にファイルオブジェクトを渡すと self._filePassed = 1 となる。この変数は __init__()でのみ定義され、変更されない。
zipfile.ZipFile の第一引数にファイルオブジェクトを渡すと self.fp = file 、つまり self.fp にBytesIOインスタンスが入る。
zipfile.ZipFile の第一引数にファイルオブジェクトを渡すと self._fileRefCnt = 1 となる。今回のスニペットだと self._fileRefCnt = 1 は1のままである。

次にソースコードにある close() メソッドを確認する。

    def close(self):
        """Close the file, and for mode 'w', 'x' and 'a' write the ending
        records."""
        if self.fp is None:
            return
        try:
            ...
        finally:
            fp = self.fp
            self.fp = None
            self._fpclose(fp)

ここでわかることは以下の点である。

内部変数fpにself.fp 、つまり BytesIOインスタンスが入る。
その後、 self.fp は None となる。
self._fpclose(fp) が実行される。

そして、ソースコードにある _fpclose()メソッドを確認する。

    def _fpclose(self, fp):
        assert self._fileRefCnt > 0
        self._fileRefCnt -= 1
        if not self._fileRefCnt and not self._filePassed:
            fp.close()

self._fileRefCnt = 1 なので、if 文に到達する際に self._fileRefCnt == 0 となり、条件式の左辺は True である。
self._filePassed = 1 なので、条件式の右辺は False である。
よって、 if not self._fileRefCnt and not self._filePassed: は False となり、 fp.close() は実行されない。

以上より、「 io.BytesIO(content) は閉じられるのだろうか」という疑問は「閉じられません」となる。

では、閉じるにはどうすればよいのか、自分で明示的に閉じるしかない。

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

こうすれば、 names: list[str] = zf.namelist() が実行された後、zipfile.ZipFileが閉じられ、 io.BytesIO が閉じられる。これで安心して眠りにつくことができる。

ガベージコレクションの存在¶

しかし、このスニペット程度ならば、ガベージコレクション先生が気を利かしてくれるはずだ。

import io
import gc
import urllib.request
import zipfile

gc.enable()

URL = "http://localhost:8000/python-3.12.2-embed-amd64.zip"

with urllib.request.urlopen(URL) as f:
    content: bytes = f.read()

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

for name in names:
    print(name)

print(gc.is_tracked(data))

...
True

当然、 io.BytesIO(content) はガベージコレクタによって捕捉されており、必要に応じて回収されるだろう。

さて、Pythonの with 文は複数のコンテキストマネージャを扱える。よって、先ほどの「io.BytesIOを明示的に閉じる」スニペットは次のようにも書ける。

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 を返す点に注意しよう。

今後の課題¶

今回のスニペット程度ではガベージコレクタによって回収されるので、閉じ忘れても特に支障はないと思われる。「ZIPファイルを手元で加工する」の内容如何では、ガベージコレクションによって回収されない状況が起き得るか、が自分の課題である。また、 gc モジュールの使い方がイマイチわかっていないので、それの調査も課題である。

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