Zenn
👍

tempfileはいいぞ

2022/05/06に公開
AWS
Python
tech

こんにちわ alivelimb です。
本記事では一時ファイル・ディレクトリ作成ができる標準パッケージtempfileを紹介します。あまり使う機会は多くないですが、知っておくと便利な時があるので是非ご一読ください。

また、tempfile を使った具体例としてAWS S3上のオブジェクトをローカルのファイルのように扱えるクラスを実装してみました。

tempfile って何ができるの

永続化する必要はないけど、一時的にファイルやディレクトリを作成できます。一時ディレクトリ(Linux であれば/tmp)に作成されるため OS による差異を気にせずに一時ファイルを作成することが可能です。

with TemporaryFile("w") as tmpf:
    tmpf.write("Hello World")

余談ですが、OS による差異なくファイルパスを扱いたい場合はpathlibもおすすめです。こちらについては記事も記事を書いているので、適宜参照してください。(pathlib はいいぞ)

TemporaryFile と NamedTemporaryFile の違い

tempfileにはいくつかのクラスが用意されていますが、TemporaryFileとNamedTemporaryFileについて紹介しておきます。違いをざっくりいうと「ユーザからそのファイルが見えるかどうか」です。

TemporaryFile

TemporaryFile は一時ファイルを作成しますが、ユーザがそのファイルパスを取得することで出来ません。そのため、一時ファイルに書き込んだ内容を取得したい場合は.seekなどを使ってカーソル位置を移動して読む必要があります。

with TemporaryFile("w+") as tmpf:
    tmpf.write("Hello World")
    tmpf.seek(0)
    print(tmpf.read())

「永続化したいわけではないけど、with ブロックを抜けたら使えなくなる一時的なファイルだとライフサイクルが短いなー」と思うケースもあるでしょう。そんな時は NamedTemporaryFile を使うと良いでしょう。

NamedTemporaryFile

NamedTemporaryFile も TemporaryFile と同様ですが、ユーザがファイルパスを取得できる点でことなります。ファイルパスは.nameで取得できます。

with NamedTemporaryFile() as tmpf:
    print(tmpf.name)

MacOS で検証した際は/varディレクトリ下に一時ファイルが作成されていました。ちなみに、TemporaryFileでも同様に.nameを取得してみると3と返ってきたので確かにファイルパスは見えないようです。

また、NamedTemporaryFile はdelete=Falseにすることで with ブロックを抜けても削除しないように設定可能です。この場合は.closeとos.unlink(filename)を後処理として実行する必要があります(参考)。

具体例: AWS S3 連携を NamedTemporaryFile で書いてみる

NamedTemporaryFile を使って AWS S3 上のオブジェクトをローカルファイルのように扱えるS3Connectorを書いてみました。コードの全量はgistで公開しているため、「説明不要、コードを見せろ」という方はこちらを参照して下さい。なお、NamedTemporaryFile を使ったサンプルコードに過ぎないので、実際に使う場合はs3fsなどを使った方が良いでしょう。 共有ファイルシステムが必要な場合は EFS などを使った方が良いでしょう。

2022.05.06 追記
AWS 公式回答で s3fs 等を用いて S3 を EC2 にマウントするのは非推奨となっていました。

s3fs などのツールを利用して S3 をマウントすることは安定性やコストの観点から非推奨としています。共有ファイルシステムが必要であれば EFS の利用をご検討ください。

使用例

bucket = "s3-bucket-name" # WRITE ME
key = "path/to/sample.txt" # WRITE ME

s3_conn = S3Connector(bucket)

with s3_conn.open(key) as f:
    print(f.read())

with s3_conn.open(key, "w") as f:
    f.write("Hello World\n")

with s3_conn.open(key, "a") as f:
    f.write("Goodbye\n")

__init__

__init__で S3 バケットを指定しています。head_bucketでバケットの存在有無を確認し、バケットがない場合はエラーをそのまま返すようにしています。

class S3Connector:
    def __init__(self, bucket: str, **kwargs: Dict) -> None:
        s3 = boto3.resource("s3", **kwargs)
        try:
            s3.meta.client.head_bucket(Bucket=bucket)
        except ClientError as err:
            raise err

        self._bucket = s3.Bucket(bucket)

Python の AWS SDK であるboto3ではbotocore.exceptions.ClientErrorで様々な例外をキャッチします。今の実装だと存在しないバケットを利用したエラーなのか、AWS の認証エラーなのか判断ができないため、本来であれば例外処理をしっかりと分けて書くべきかと思います。(今回はあくまでtempfileの具体例なので手抜きをしています。。。)

open

今回は読み込み、書き込み、追記の 3 つモードを用意しました。

S3ConnectorMode = Literal["r", "w", "a"]

# 中略

@contextmanager
def open(self, s3key: str, mode: S3ConnectorMode = "r") -> Generator[IO, None, None]:
    if mode == "r":
        yield from self._read(s3key)
    elif mode == "w":
        yield from self._write(s3key)
    elif mode == "a":
        yield from self._add(s3key)
    else:
        raise Exception("invalid mode")

contextlibを用いることでwithブロックと組み合わせる関数が書きやすくなります。

_read

ようやくNamedTemporaryFileの出番です。まずは読み込みです。前処理として S3 から対象オブジェクトをローカルにダウンロードします。この時、ダウンロード先をNamedTemporaryFileで作成した一時ファイルに設定しています。

def _read(self, s3key: str) -> Generator[IO, None, None]:
    # 存在確認
    try:
        self._bucket.meta.client.head_object(Bucket=self._bucket.name, Key=s3key)
    except ClientError as err:
        raise err

    # S3のオブジェクトをローカルに一時ファイルとして作成
    with NamedTemporaryFile(delete=False, mode="w") as f_tmp:
        tmpfile_name = f_tmp.name
        self._bucket.download_file(s3key, tmpfile_name)

    # withブロックに読み込み専用ファイルIOを渡す
    # withブロック終了後に一時ファイルを閉じて削除する
    f_yield = open(tmpfile_name)
    try:
        yield f_yield
    finally:
        f_yield.close()
        os.unlink(tmpfile_name)

コンテキストマネージャを書くときの注意点として、呼び出し元のwithブロック内で例外が発生した時に備えておく必要があります。これはtry-finallyで囲むことで例外が発生した場合でも安全に後処理を行うことができます。

_write

次に書き込みです。前処理として書き込み用の一時ファイルを作成しておき、後処理としてファイルを S3 にアップロードします。

def _write(self, s3key: str) -> Generator[IO, None, None]:
    # 書き込み用の一時ファイルを作成
    f_yield = NamedTemporaryFile(delete=False, mode="w")
    tmpfile_name = f_yield.name

    # withブロックに書き込み専用ファイルIOを渡す
    # withブロック内で例外が発生した場合は一時ファイルを閉じて削除する
    try:
        yield f_yield
    except Exception as err:
        f_yield.close()
        os.unlink(tmpfile_name)
        raise err

    # withブロックが正常終了した場合
    # S3にアップロードしてから一時ファイルを閉じて削除する
    try:
        f_yield.flush()
        self._bucket.upload_file(tmpfile_name, s3key)
    except ClientError as err:
        raise err
    finally:
        f_yield.close()
        os.unlink(tmpfile_name)

ファイルに書き込む際は、書き込んだ内容がバッファに溜まっているだけで、ファイルに書き込めていない可能性があります。これを回避するために.flushを呼び出すことで、バッファの内容をファイルに書き込んでからアップロードします。

_add

最後は追記です。これまでに読み込みと書き込みの組み合わせ技になっているので、特に説明は不要かと思います。

def _add(self, s3key: str) -> Generator[IO, None, None]:
    # 存在確認
    try:
        self._bucket.meta.client.head_object(Bucket=self._bucket.name, Key=s3key)
    except ClientError as err:
        raise err

    # S3のオブジェクトをローカルに一時ファイルとして作成
    with NamedTemporaryFile(delete=False, mode="w") as f_tmp:
        tmpfile_name = f_tmp.name
        self._bucket.download_file(s3key, tmpfile_name)

    # withブロックに追記専用ファイルIOを渡す
    # withブロック終了後に一時ファイルを閉じて削除する
    f_yield = open(tmpfile_name, "a")
    try:
        yield f_yield
    except Exception as err:
        f_yield.close()
        os.unlink(tmpfile_name)
        raise err

    # withブロックが正常終了した場合
    # S3にアップロードしてから一時ファイルを閉じて削除する
    try:
        f_yield.flush()
        self._bucket.upload_file(tmpfile_name, s3key)
    except ClientError as err:
        raise err
    finally:
        f_yield.close()
        os.unlink(tmpfile_name)

まとめ

本記事では一時ファイル・ディレクトリ作成ができる標準パッケージtempfileを紹介しました。正直そんなに使う機会はないのですが、覚えておいて損はない標準パッケージだと思います。

Discussion