状態とふるまいを持つモデルを実装する (3) - デコレータによる実装例

前回 （状態とふるまいを持つモデルを実装する (2)）は、３つの状態を持つ StopWatch を、二つの方法で実装した。今回は StopWatch を decorator を利用して実装する。

まずは、状態遷移を decorator で記述する。状態遷移するメソッドかわかりやすくなる。

from enum import auto


def transit(state):
    def decorator(func):
        def inner(self, *args, **kwargs):
            self.start_stop, self.reset = self._TRANSIT[state]
            func(self, *args, **kwargs)
        return inner
    return decorator


class StopWatch:
    WAIT, PAUSE, MEASURE = auto(), auto(), auto()

    def __init__(self):
        self._TRANSIT = {StopWatch.WAIT: (self.start, lambda *args: None),
                         StopWatch.PAUSE: (self.start, self.reset_time),
                         StopWatch.MEASURE: (self.pause, lambda *args: None)}
        self._TRANSIT_REVERSED = {v: k for k, v in self._TRANSIT.items()}
        self.start_stop, self.reset = self._TRANSIT[StopWatch.WAIT]

    @property
    def state(self):
        return self._TRANSIT_REVERSED[self.start_stop, self.reset]

    @transit(MEASURE)
    def start(self):
        pass

    @transit(PAUSE)
    def pause(self):
        pass

    @transit(WAIT)
    def reset_time(self):
        pass

さらに状態遷移表をデコレータで切り出してみる。behavior デコレータはある状態のときのふるまいを表す。

from enum import auto
from state_machine import behavior, transit


class StopWatch:
    WAIT, MEASURE, PAUSE = auto(), auto(), auto()

    def __init__(self):
        self.state = StopWatch.WAIT

    @behavior(WAIT, "start")
    @behavior(MEASURE, "pause")
    @behavior(PAUSE, "start")
    def start_stop(self):
        pass

    @behavior(PAUSE, "reset_time")
    def reset(self):
        pass

    @transit(MEASURE)
    def start(self):
        pass

    @transit(PAUSE)
    def pause(self):
        pass

    @transit(WAIT)
    def reset_time(self):
        pass

state_machine.py

from functools import wraps


def transit(state):
    def decorator(func):
        @wraps(func)
        def inner(self, *args, **kwargs):
            self.state = state
            func(self, *args, **kwargs)
        return inner
    return decorator


def behavior(state, function):
    def decorator(func):
        @wraps(func)
        def inner(self, *args, **kwargs):
            if self.state == state:
                getattr(self, function)(*args, **kwargs)
            else:
                func(self, *args, **kwargs)
        return inner
    return decorator

状態を self.state に保持する前提ではあるが、他のモデルでも同じ decorator が使用できるため、モジュールstate_machineとして切り出した。

デコレータを処理するタイミングでは、まだ関数がアンバウンドなので、ふるまいを文字列で与えている。関数の__name__ を利用して実行時にバウンドメソッドを呼ぶこともできる。（該当箇所のみ以下に示す。）
なお、functools.wrap() を使用しない decorator があると想定とおり動作しない。（functools.wrap() を使用すると decorator が関数の __name__ を書き換えない。）

state_machine.py

def behavior(state, function):
...
            if self.state == state:
                getattr(self, function.__name__)(*args, **kwargs)
...

...
    @behavior(PAUSE, reset_time)
    def reset(self):
        pass
...