Enum(Python)の小ネタ

PythonのEnumでCallableな値を持ちたいと思ったことありませんか？(強引)
そして実装してみるととある問題に直面しませんか？

あれ…キー指定して呼び出そうとするとKeyError発生しないか…？

つまりはこういうことです。

import enum
import sys
from functools import partial
from typing import List, Callable

class Color(enum.Enum):
    RED: Callable[[str], str] = lambda m: f"<font color=\"red\">{m}</font>"
    BLUE: Callable[[str], str] = lambda m: f"<font color=\"blue\">{m}</font>"
    GREEN: Callable[[str], str] = lambda m: f"<font color=\"green\">{m}</font>"

    @property
    def value(self) -> Callable[[str], str]:
        return super().value

def main(argv: List[str]) -> int:
    color, msg = argv

    c = Color[color]
    print(c.value(msg))

    return 0

if __name__ == "__main__":
    main(sys.argv[1:])

ColorクラスのRED, BLUE, GREENの値がlambdaになっています(わざわざこんな実装にしないというツッコミは無しで)
スクリプトの第一引数に色、第二引数にメッセージを設定できるようにして、main.pyとして保存してコンソールから呼び出すと

python main.py RED 'Hello, World!'
Traceback (most recent call last):
  File "/Users/k/workspace/PythonProject/enum_test/main.py", line 24, in <module>
    main(sys.argv[1:])
  File "/Users/k/workspace/PythonProject/enum_test/main.py", line 18, in main
    c = Color[color]
  File "/Users/k/.pyenv/versions/3.9.4/lib/python3.9/enum.py", line 408, in __getitem__
    return cls._member_map_[name]
KeyError: 'RED'

何故かKeyErrorが発生します。

回避策は色々あるようですが、functools#partialを使うのが楽な方法そうです。
先程のEnumクラスの値のlambda部分をpartialで包んでpartialオブジェクトとしてあげます。

import enum
import sys
from functools import partial
from typing import List, Callable

class Color(enum.Enum):
    # partial
    RED: Callable[[str], str] = partial(lambda m: f"<font color=\"red\">{m}</font>")
    BLUE: Callable[[str], str] = partial(lambda m: f"<font color=\"blue\">{m}</font>")
    GREEN: Callable[[str], str] = partial(lambda m: f"<font color=\"green\">{m}</font>")

    @property
    def value(self) -> Callable[[str], str]:
        return super().value

def main(argv: List[str]) -> int:
    color, msg = argv

    c = Color[color]
    print(c.value(msg))

    return 0

if __name__ == "__main__":
    main(sys.argv[1:])

上記コードを先程と同様に呼び出すと、

python main.py RED 'Hello, World!'
<font color="red">Hello, World!</font>

うまくいきました！
なぜかEnumの値に直接functionを持たせるとエラーになってしまうので、代理としてfunctionを包んだオブジェクトを指定してあげると良いようです。
(ここらへんの詳しい原因調査は行っていないのでエラーになる理由が分かる方いれば教えて頂けると助かります…)