PythonのEnumを関数と紐付け、Callableにする

いくつかの選択肢からアルゴリズムを選ばせるような時、列挙型の各値に関数を紐づけたいことがある。

Javaなら列挙型の各値にメソッドを設定できるのだが、Pythonではできない。Enum自体をCallableにする方向性で、なんとかキレイな書き方がないか考えてみる。

値はどうでもいい場合

PythonのEnumは値を自由に設定できる。だから値そのものを関数にしてしまえばいい。

だがEnumを継承したクラスの中で定義された関数はメソッドと判断され、列挙される値に含まれなくなってしまう。

class Dame(Enum):
    # 値を定義したことにならず、メソッドを定義したことになる
    A = lambda: print("A")

list(Dame)
# 空リスト

そこで、このstackoverflowのアイデアを借りて、関数が含まれるタプルを値に使うことにする。

from enum import Enum

class CallableEnum(Enum):
    """
    このクラスを継承してCallableなEnumを作る
    """
    def __call__(self, *args, **kwargs):
        return self.value[0](*args, **kwargs)

def register(func):
    """
    関数が含まれるタプルを簡単に定義するためのデコレータ
    """
    return func,

使い方はこう。

class Test(CallableEnum):
    @register
    def A():
        print("A")

    @register
    def B():
        print("B")

Test.A()
# A

値を自由に設定したい場合

Enumを継承したクラスTに、Tから関数への辞書を持たせて、__call__の中で参照する。

ただし「Tから関数への辞書」をクラス定義中にTの属性として持たせてしまうと、それも列挙型の値の一つになってしまう。Tがクラスとして出来上がった後で付け加えないといけない。

したがって、クラス定義が終わった後に関数を付け加える書き方になる。

from enum import Enum

class Dnum(Enum):
    """
    Dispatching Enum。こいつを継承する
    """
    @classmethod
    def register(cls, key):
        if not hasattr(cls, "table"):
            cls.table = {}
        def registration(func):
            cls.table[key] = func
        return registration

    def __call__(self, *args, **kwargs):
        return self.__class__.table[self](*args, **kwargs)


def register(enum):
    def ret(func):
        enum.__class__.register(enum)(func)
    return ret

こんな感じで使う。

from enum import auto

class Test(Dnum):
    A = "A"
    B = "B"

@register(Test.A)
def _A():
    print("A")

@register(Test.B)
def _B():
    print("B")

Test.A()
# A

値を自由に設定しつつ、関数をクラス定義の中に書く裏技ないかな……