Python で一時的に関数出力を凍結 (キャッシュ) する

TL;DR 🎄

これはアドベントカレンダーの15日目の記事です。 今日は、Python の関数出力を一時的に変化させないようにキャッシュするという、わりとマニアックな話のメモです。

Temporary caching of function’s return

例えば、ユーザからクエリを受け取り、何かファイルからデータを読み込み、該当する値を返す処理を考えます。

def get_value(query):
    data = load(filename)
    return data[query]

このように書くことで、ファイルに変更があった場合もその都度関数がロードするため、常に最新の値を得ることができます。 一方、以下のように大量のクエリを処理する場合、クエリの数だけロードが発生するので、もしファイルが巨大な場合I/Oで非常に時間がかかることが予想されます。

for query in million_query_list:
    value = get_value(query)
    ... # 値に対する何かの処理

そこで、get_value 関数自体は変更することなく、for 文の中だけ load 関数の出力を凍結 (キャッシュ) することができないか考えてみます。 これは、load 関数をキャッシュ化した関数を作成し、一時的に load 関数と置き換えてあげることで可能になります。 また、”for 文の中” などのコンテクストを扱いたいので、with 文を使ったコンテキストマネージャを使ってあげれば良いことも分かります。

Freeze context manager

というわけで、以下のようなコンテクストマネージャを作成しました。 with freeze(func, module) のように関数とこれが属するモジュールを与えることにより、with 文の中では func() が最初に返した値を常に返すようになります。

from inspect import getmodule

class freeze:
    def __init__(self, func, module=None):
        self.func = func
        self.module = module or getmodule(func)

    def frozen(self, *args, **kwargs):
        if hasattr(self, 'cache'):
            return self.cache

        self.cache = self.func(*args, **kwargs)
        return self.cache

    def __enter__(self):
        setattr(self.module, self.func.__name__, self.frozen)

    def __exit__(self, exc_type, exc_value, traceback):
        setattr(self.module, self.func.__name__, self.func)

An example

これをテストするため、以下のような現在時刻を返す now 関数を作成しました。

from datetime import datetime

def now():
    return datetime.now()

これを普通に3連続で実行すると、マイクロ秒単位で異なる時間を返すはずです。

print(now())
print(now())
print(now())

# 2018-12-15 23:08:20.573039
# 2018-12-15 23:08:20.573235
# 2018-12-15 23:08:20.573312

今度はこれを、freeze コンテクストマネージャの下で実行してみます。

with freeze(now):
    print(now())
    print(now())
    print(now())

print(now())

# 2018-12-17 23:08:29.831403
# 2018-12-17 23:08:29.831403
# 2018-12-17 23:08:29.831403
# 2018-12-17 23:08:29.831526

このように、with 文中では、最初に呼び出された時刻を返し続けることが分かります。 また、with 文の外では、再び現在時刻を返していることが確認できました。