Pythonで一時ディレクトリを作って安全に後始末する

前置き：一時ファイルの作成

処理の途中で一時ファイルを作っていろいろ作業して、最後に削除したいということがありますよね。
こういう場合はTemporaryFileとかNamedTemporaryFileを使います。

Python2.7

import tempfile
with tempfile.TemporaryFile() as f:
    print >>f, "Test"
    f.seek(0)
    print f.read()

Python3.6

import tempfile
with tempfile.TemporaryFile("w+") as f:
    print("Test", file=f)
    f.seek(0)
    print(f.read())

Python 3だと一時ディレクトリも簡単

使うファイルが多いとか、外部ツールにファイルを作らせるとかで、一時ファイルではなくて一時ディレクトリを作りたいということがあるかもしれません。

Python 3（3.2以降）にはTemporaryDirectoryというものがあって、これを使うと一時ディレクトリを作ったうえで、使い終わったディレクトリを自動的に後始末してくれます。

Python3.6

import tempfile
import os
with tempfile.TemporaryDirectory() as dname:
    print(dname)                 # /tmp/tmpl2cvqpq5
    print(os.path.exists(dname)) # True
    with open(os.path.join(dname, "test.txt"), "w") as f:
        print("test", file=f)
print(os.path.exists(dname))     # False

with構文に指定した場合、たとえwith構文の中で例外が起きたとしても先に後始末してから死にます。
with構文をtry - except - finallyで囲んだ場合でも、例外が起きた時点でまず後始末されます。便利ですね。

しかしPython 2にはTemporaryDirectoryがありません。
mkdtempはありますが、自動で後始末はしてくれません。ただ一時ディレクトリを作ってくれるだけ。
ということで、Python 2でも後始末させましょう。

Python 2で一時ディレクトリ

要件

• インスタンスが生成されたらmkdtempで一時ディレクトリを作成し、cleanupメソッドが呼ばれたら削除する
• インスタンスのname属性から一時ディレクトリ名を取得する
• with構文に指定した場合、asの後ろに指定した変数にはディレクトリ名（TemporaryDirectory.name相当）が入る。with構文から抜けたら自動的にcleanupメソッドを呼ぶ

Python 3の実装をベースに作ってみた

Python 3.6のライブラリの中から必要な部分を抜き出し、__init__のデフォルト引数をPython 2のmkdtempに寄せています。
その他、name属性を読み取り専用にするとか、ちょっと手を加えています。
（Linux環境なら、元ネタは /usr/lib/python3.6/tempfile.py あたりにあります）

Python2.7

import tempfile
import shutil
import os

class TemporaryDirectory(object):
    def __init__(self, suffix="", prefix="tmp", dir=None):
        self.__name = tempfile.mkdtemp(suffix, prefix, dir)

    def __enter__(self):
        return self.__name

    def __exit__(self, exc, value, tb):
        self.cleanup()

    @property
    def name(self):
        return self.__name

    def cleanup(self):
        shutil.rmtree(self.__name)

with TemporaryDirectory() as dname:
    print dname                 # /tmp/tmpslYr3Q
    print os.path.exists(dname) # True
    with open(os.path.join(dname, "test.txt"), "w") as f:
        print >>f, "test"
print os.path.exists(dname)     # False

実行結果はこんな感じ。

/tmp/tmpslYr3Q
True
False

使用例をもう一つ。例外を発生させてみます。

Python2.7

d = TemporaryDirectory()
print os.path.exists(d.name)
try:
    with d:
        print d.name
        d.name = "/tmp/123456"
finally:
    print d.name
    print os.path.exists(d.name)

実行結果はこうなります。（IPythonの場合）

True
/tmp/tmpJcC_mz
/tmp/tmpJcC_mz
False
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-17-372b6ae08491> in <module>()
      4     with d:
      5         print d.name
----> 6         d.name = "/tmp/123456"
      7 finally:
      8     print os.path.exists(dname)

AttributeError: can't set attribute

このように、nameに値を設定しようとすると例外が発生するようになっています。
また、例外が発生したときにfinallyのprint文が実行されますが、os.path.existsの結果がFalseになっていることから、finally節が実行される前に後始末が行われていることが分かります。