built-in object を拡張する禁断の果実を齧ろう

Python は元々 monkey-patching しやすい言語ですが、C で書かれた組み込み型を拡張しようとすると TypeError が起きます。

>>> str.foo = lambda: 'foo'
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: can't set attributes of built-in/extension type 'str'

そんなときは forbiddenfruit を使うと、組み込み型も拡張できます。

pip でインストール

pip install forbiddenfruit

datetime, time に monkey-patch

実際に現在日時を返す datetime built-in object のクラスメソッドを翌日の同時刻を返すメソッドに拡張してみました。また、time モジュールの built-in function も拡張しました。こちらは forbiddenfruit は不要で、単純にすり替えるだけです。

hook_datetime.py

import datetime, time
from forbiddenfruit import curse
original_now = datetime.datetime.now
original_utcnow = datetime.datetime.utcnow
original_today = datetime.date.today

def tomorrow_now(self, tz=None):
    now = original_now(tz)
    return now + datetime.timedelta(days=1)

def tomorrow_utcnow(self):
    now = original_utcnow()
    return now + datetime.timedelta(days=1)

def tomorrow_today(self):
    return tomorrow_now(datetime.datetime).date()

def patch_tomorrow_time():
    global time
    original_localtime = time.localtime
    def _localtime(secs=None):
        if not secs:
            return tomorrow_now(datetime.datetime).timetuple()
        else:
            return original_localtime(secs)
    time.localtime = _localtime
    def _time():
        return time.mktime(time.localtime())
    time.time = _time
    original_gmtime = time.gmtime
    def _gmtime(secs=None):
        if not secs:
            return tomorrow_utcnow(datetime.datetime).timetuple()
        else:
            return original_gmtime(secs)
    original_asctime = time.asctime
    def _asctime(t=None):
        if not t:
            return original_asctime(tomorrow_now(datetime.datetime).timetuple())
        else:
            return original_asctime(t)
    time.asctime = _asctime
    def _ctime(secs=None):
        return time.asctime(time.localtime(secs))
    time.ctime = _ctime

curse(datetime.datetime, 'now', classmethod(tomorrow_now))
curse(datetime.datetime, 'utcnow', classmethod(tomorrow_utcnow))
curse(datetime.date, 'today', classmethod(tomorrow_today))
patch_tomorrow_time()

built-in function のパッチの方が長くて汚くなってしまった。。
しかしまあ、これであっさりと datetime.now() などが翌日を返すようになりました。Django の ./manage.py runserver なんかも翌日に進んだ状態で普通に動きました。

雑感

テストやデバッグで時間を変えたいケースがあったりするのでシステム時間変えるのは腰が重いなーと感じて試しに書いてみました。

datetime の例は monkey-patching するメソッドに漏れがありそうな気がしてますが、production では決して使わないコードなんで出たとこ勝負で直していけば良いかなと思います。

forbiddenfruit 自体は 150 行前後の短いコードなので読んでみると ctypes について勉強になる感じでした。