Python でスタックを消費しない末尾再帰

末尾再帰でスタックを消費しないようにするためのデコレータが紹介されています。

New Tail Recursion Decorator « Python recipes « ActiveState Code

Pythonで末尾再帰最適化をする。

Pythonのクロージャで末尾再帰最適化をする。

上記の各リンク先のコードが書かれてから3年以上経過しています。 その間に Python 3 がメジャーなってきました。 そこで、 Python 3 から導入されているnonlocal文を使い、一番下のリンク先のコードを書き直してみました。

from functools import wraps

def recur(func):
    _func = func
    _firstcall = True
    _CONTINUE = object()
    _args = None

    @wraps(func)
    def _recur(*args, **kwargs):
        nonlocal _func, _firstcall, _CONTINUE, _args
        if _firstcall is True:
            func = _func
            CONTINUE = _CONTINUE
            _firstcall = False
            try:
                while True:
                    result = func(*args, **kwargs)
                    if result is CONTINUE:
                        args, kwargs = _args
                    else:
                        return result
            finally:
                _firstcall = True
        else:
            _args = args, kwargs
            return _CONTINUE
    return _recur

変数への再代入が気持ち悪いと感じる方もいるかもしれませんが、nonlocal文のおかげで大分スッキリしたコードになったと思います。

誤りを見つけた方は是非教えてください。