【Python】Enum を拡張して便利に使いたい

Python で Enum を拡張して便利に使う方法について説明します。

自分的には区分値などを定義するときに重宝しています。

拡張 Enum

以下のように Enum を継承したクラスを作成します。

from enum import Enum

class TypeBase(Enum):
    def __init__(self, val, lavel, data: dict = None):
        self.val = val
        self.label = label
        self.data = data

    @classmethod
    def get_members(cls):
        """全メンバを取得"""
        return [*cls.__members__.values()]

    @classmethod
    def get_by_val(cls, val):
        """値からメンバを取得"""
        for c in cls.get_members():
            if c.val == val:
                return c
        return None

    @classmethod
    def get_choices(cls, empty_label=None):
        """選択肢用"""
        choices = []
        if empty_label:
            choices.append(("", empty_label))
        choices.extend([c.value for c in cls])
        return choices

# その他便利なメソッドを定義

使い方

拡張 Enum を使って定義をします。

class ExampleType(TypeBase):
    ONE = (1, "one")
    TWO = (2, "two", {"option": "ext"})

# 使うとき
ExampleType.ONE.val
# 1

example = ExampleType.get_by_val(2)
# ExampleType.TWO

使いやすいようにメソッドを追加したりできるのでかなり便利です。

参考にしてみてください。