Python Tips: Enum 型を使いたい

Python で Enum 型(列挙型)を使う方法について、手短に説明してみます。

お断り: 以下に記載するコードについては動作確認はしていますが、私はたくさん Enum 型を使ってきたわけでは無いので、理解が間違っている可能性もあります。ご了承ください。

Python の Enum 型

Python では Python 3.4 から enum という標準モジュールが追加され、 Python 本体に Enum 型が同梱されるようになりました。

import enum

Enum 型の基本的な使い方

定義

オリジナルの Enum 型は enum.Enum 型を継承して作成します。

from enum import Enum

class Status(Enum):
    ACTIVE = 1
    INACTIVE = 0
    CANCELED = -1

アイテム(=インスタンス)はクラスアトリビュートとして定義します。このサンプルの場合は、 Status.ACTIVE Status.INACTIVE Status.CANCELED という 3 つのアイテムを定義していることになります。

ふるまい

Enum 型のふるまいを unittest で確認してみましょう。

from enum import Enum
from collections import OrderedDict
import unittest

class Status(Enum):
    ACTIVE = 1
    INACTIVE = 0
    CANCELED = -1

class TestStatus(unittest.TestCase):
    def test_types(self):
        # クラスのアトリビュートは自動的にインスタンスとして扱われる
        self.assertIsInstance(Status.ACTIVE, Status)
        self.assertIsInstance(Status.INACTIVE, Status)
        self.assertIsInstance(Status.CANCELED, Status)

    def test_for(self):
        # iterable プロトコルを持っているので for ループで回せる
        it = iter(Status)
        self.assertIs(next(it), Status.ACTIVE)
        self.assertIs(next(it), Status.INACTIVE)
        self.assertIs(next(it), Status.CANCELED)
        with self.assertRaises(StopIteration):
            next(it)

    def test_in(self):
        # iterable プロトコルを持っているので in 演算子も使える
        self.assertIn(Status.CANCELED, Status)

    def test_name_and_value(self):
        # キーと値はアトリビュート `name` と `value` で取得できる
        self.assertIs(Status.INACTIVE.name, 'INACTIVE')
        self.assertIs(Status.INACTIVE.value, 0)

    def test_accessors(self):
        # 各アイテムにアクセスする方法として次の 3 種類が要されてている
        self.assertTrue(Status['CANCELED'] == Status(-1) == Status.CANCELED)

    def test_inequality_with_value(self):
        # キーや値と勝手に等しくなったりはしない
        self.assertNotEqual(Status.CANCELED, 'CANCELED')
        self.assertNotEqual(Status.CANCELED, -1)

    def test_dunder_members(self):
        # クラスの __members__ アトリビュートには OrderedDict が入っている
        self.assertEqual(
            Status.__members__,
            OrderedDict(
                {
                    'ACTIVE': Status.ACTIVE,
                    'INACTIVE': Status.INACTIVE,
                    'CANCELED': Status.CANCELED,
                }
            ),
        )

なんだか直感的なふるまいをしてくれる印象がありますが、いかがでしょうか。 enum.Enum を使って定義できる Enum 型は私が Enum 型に期待する機能をひととおり揃えていてかつシンプルなので、私にはとても使いやすい印象です。

自動採番

Python 3.6 以降なら、 enum.auto() を使って各アイテムの値を自動でセットすることができます。

from enum import Enum, auto
import unittest

class Prefecture(Enum):
    HOKKAIDO = auto()
    AOMORI = auto()

class TestPrefecture(unittest.TestCase):
    def test_auto_values(self):
        # auto() を使うと 1 始まりで整数が自動的に振られる
        self.assertEqual(Prefecture.HOKKAIDO.value, 1)
        self.assertEqual(Prefecture.AOMORI.value, 2)

シンプル定義

Enum 型の定義方法としてメジャーなのはおそらく上の enum.Enum クラスを継承する方法ですが、次のように enum.Enum を関数のように使用してワンラインで Enum 型を定義する方法も用意されています。

PublishedStatus = Enum('PublishedStatus', 'PUBLISHED DRAFT')

class TestPublishedStatus(unittest.TestCase):
    def test_instances(self):
        self.assertIsInstance(PublishedStatus.PUBLISHED, PublishedStatus)
        self.assertIsInstance(PublishedStatus.DRAFT, PublishedStatus)

    def test_values(self):
        self.assertEqual(PublishedStatus.PUBLISHED.value, 1)
        self.assertEqual(PublishedStatus.DRAFT.value, 2)

この形で Enum 型を定義する場合は、 enum.Enum の第 1 引数に Enum 型の名前を、第 2 引数にアイテムをスペースや , で区切りの文字列で渡します。第 2 引数は listdict で渡すことも可能です。詳細は公式ページの Functional API の節でご確認ください。

というわけで、 Python の Enum 型の使い方についてでした。

シンプルな用途で Enum 型を使う場合は以上の内容を知っていれば十分ではないかと思いますが、その他細かな機能がもろもろ用意されているので、詳細を知りたい方は公式ページをご覧になってから使うのがよいかと思います。

参考