Pythonの列挙型 (enum)

Python

これまで定数クラスを作ってやり過ごすことが多かったのですが、enumが比較的使いやすかったので、

enum

enumはPython 3.4から標準ライブラリに追加されたモジュールで、列挙型をサポートするものです。

docs.python.org

実装

最初にenumをインポートしておきます。Python 3.7.4を使っています。

import enum

BTreeなどの木構造のノードを考えてみます。各ノードは根ノード (ROOT) 、中間ノード (INTERNAL) 、葉ノード (LEAF) の3種類のいずれかに分類されるものとします。

それをTreeNodeTypeという列挙型で表現すると、以下のようになります。Enumクラスを継承することで実装します。

  • 列挙子 (ROOTなどの識別子) はnameとvalueとプロパティを持ってます
  • 列挙子同士は比較可能ですが、列挙子と同じEnumオブジェクトではない値との比較はFalseになります (最終行のFalse)
class TreeNodeType(enum.Enum):
    ROOT = 1
    INTERNAL = 2
    LEAF = 3

print(type(TreeNodeType.ROOT))  # <enum 'TreeNodeType'>
print(TreeNodeType.ROOT)  # TreeNodeType.ROOT

print(TreeNodeType.ROOT.name)  # ROOT
print(TreeNodeType.ROOT.value)  # 1

root1 = TreeNodeType.ROOT
print(root1 == TreeNodeType.ROOT)  # True

root2 = TreeNodeType.ROOT
print(root1 == root2)  # True

print(root1 == 1)  # False

列挙子の値を自分で埋めるのがめんどくさい、というのであれば、auto関数 (3.6から追加) を使うと1開始の値でセットしてくれます。

class TreeNodeType(enum.Enum):
    ROOT = enum.auto()
    INTERNAL = enum.auto()
    LEAF = enum.auto()

print(list(TreeNodeType))
# [<TreeNodeType.ROOT: 1>, <TreeNodeType.INTERNAL: 2>, <TreeNodeType.LEAF: 3>]

デフォルトでは列挙子の値は重複してもOKです。

class TreeNodeType(enum.Enum):
    ROOT = 1
    INTERNAL = 1  # ROOTと同じ
    LEAF = 3

print(list(TreeNodeType))
# [<TreeNodeType.ROOT: 1>, <TreeNodeType.LEAF: 3>]
print(TreeNodeType.ROOT == TreeNodeType.INTERNAL)  # True

ですがこれは一部のマイナーケースを除いて事故を招くだけです。uniqueアノテーションをEnumクラスに付与することで、定義時に一意に定まらない場合は例外をスローするようにできます。

@enum.unique
class TreeNodeType(enum.Enum):
    ROOT = 1
    INTERNAL = 1
    LEAF = 3
# -> ValueError: duplicate values found in <enum 'TreeNodeType'>: INTERNAL -> ROOT

__members__プロパティを使うとOrderedDictで列挙子と値を取り出すこともできます。

class TreeNodeType(enum.Enum):
    ROOT = 1
    INTERNAL = 2
    LEAF = 3
    
for k, v in TreeNodeType.__members__.items():
    print(k, v.value)
    
# ROOT 1
# INTERNAL 2
# LEAF 3

先程列挙子と同じEnumクラスに属さない値は比較できない、という話をしましたが、Enumの代わりにIntEnum型 (3.6から追加) を継承するとそれができるようになります。

  • それ以外はEnum型と同じです
class TreeNodeType(enum.IntEnum):
    ROOT = 1
    INTERNAL = 2
    LEAF = 3
    
print(root1 == 1)  # True

ファイルパーミッションのようなフラグを表現するEnum派生型として、FlagとIntFlagがあります。通常Enumは列挙子に1つの値しかとりませんが、

  • 和 (|) や積 (&) などのビット演算子が利用でき、計算後の値も同一のEnum型となります
  • inを使うと特定のフラグがセットされているのかどうかを検査できます
class Permission(enum.IntFlag):
    EXECUTE = 1
    WRITE = 2
    READ = 4
    
print(Permission.READ == 4)  # True
print(Permission.READ + Permission.WRITE)  # 6

print(Permission.READ | Permission.WRITE)  # Permission.READ|WRITE
print(type(Permission.READ | Permission.WRITE))  # <enum 'Permission'>
print(Permission.READ | Permission.WRITE == 6)  # True

print((Permission.READ | Permission.WRITE)  & (Permission.WRITE | Permission.EXECUTE))
# Permission.WRITE

print(Permission.READ in Permission.READ | Permission.WRITE)  # True

まとめ

Pythonの列挙型Enumについてさらっと解説しました。