🐍

【Python】Enum の値に重複がないことを保証する方法

に公開

Python の Enum は Pydantic と合わせると API の定義やバリデーションが楽に実装されてとても便利です。しかし定義の数が多くなると、定義に間違いがないか不安になります。

Visual Studio Code 用の Python 向け言語サーバー拡張機能である Pylance を使えば、Enum の名前の重複は静的解析でエラーになってくれます。

from enum import StrEnum

class Sample(StrEnum):
    A = "01"
    A = "02"  # <- エラー

一方、値の方の重複は静的解析でエラーになってくれません。(ruff はチェックしてくれるようですが、Enum ではなく StrEnum だとチェックしてくれないようです)。

class Sample(StrEnum):
    A = "01"
    B = "01"  # <- エラーにならない

Enum の要素数が少なければ気にすることではないですが、都道府県コードのように要素数が多くなると定義ミスがないか不安になるので、チェックしてくれると助かります。

この記事では、重複をチェックする方法を紹介します。

unique デコレーターを使う方法

いちばん簡単な方法は enum モジュールに用意されている unique デコレーター (公式ドキュメント) を使う方法です。

from enum import StrEnum, unique

@unique
class Sample(StrEnum):
    A = "01"
    B = "01"

これは value を参照していると思われ、以下のような特殊な定義でも判定してくれます。

@unique
class PrefCode(Enum):
    """aaa"""

    HOKKAIDO = ("01", "北海道")
    AOMORI = ("01", "青森")

    def __new__(cls, code: str, _) -> "PrefCode":
        obj = object.__new__(cls)
        obj._value_ = code
        return obj

    def __init__(self, _: str, name: str) -> None:
        self.pref_name = name

重複判定は実行時に行われ、重複があれば ValueError が発生します。簡単ですが、実行時に毎回行われるというのは微妙です。

単体テストで担保する

Enum の値に重複があることの判定自体はそれほど難しくなく、__members__ で要素一覧が取得できるので、それを使えば以下のように実装できます。__members__ は名前を key、Enum の要素を value とする辞書オブジェクトです。

def duplicate_check():
    values = [e.value for e in PrefCode.__members__.values()]
    assert len(values) == len(set(values))

ちなみに Enum 自体が Iterable なので以下のような for 分を書くこともできますが、値が重複しているとその分が欠落するようです。

for e in PrefCode:
    print(e.value)

pytest を使うと、以下のようになります。

import pytest

@pytest.mark.parametrize(
    ["enum_cls"],
    [
        pytest.param(PrefCode),
        ...  # Enum を追加していく
    ],
)
def test_duplicate(enum_cls):
    values = [e.value for e in enum_cls.__members__.values()]
    assert len(values) == len(set(values))

まとめ

正直なところ静的解析でなんとかしたいですが、その方法が見つからないならこの記事で紹介した方法が良いと思います。

GitHubで編集を提案

Discussion