mizzsugar’s blog

Pythonで学んだことや読書録を書きます。

そのif文、Enumにしてみませんか。

これは、Python Advent Calendar 17日目の記事です。

Python Advent Calendar 2019 - Qiita

この記事では、私が好きな標準モジュールのうちの一つenumモジュールの基本的な使い方と enumモジュールを使ってリファクタリングする例を紹介します。

enumモジュールとは

列挙型クラスをサポートします。

列挙型とは

Pythonのドキュメントにはこう書かれています。

列挙型は、一意の定数値に束縛された識別名 (メンバー) の集合です。列挙型の中でメンバーの同一性を比較でき、列挙型自身でイテレートが可能です。

https://docs.python.org/ja/3.8/library/enum.html

enum.Enumを使用した例

例えば、血液型を表すクラスがあるとします。 血液型はA, B, O, ABの4種類から成り立っています。 (より細かい分類がありますが、ここではこの4種類のみ使うとします)

一般的なクラスで表現するとこうなります。

class BloodType:
    Def __init__(self, type: str):
        self. type = type


ただ、上記の例だとA, B, O, AB以外の文字列が入ることを許可します。

ということで、__init__内で検知します。

class BloodType:
    Def __init__(self, type: str):
        if type not in (‘A’, ‘B’, ‘O’, ‘AB’):
            raise ValueError(‘妥当な血液型ではありません’)
        self. type = type


これも悪くないですが、if文を使わなくても 正しい値をとる方法があります。

血液型を列挙型クラスで表現することです。

import enum


class BloodType(enum.Enum):
    A = 'A'
    B = 'B'
    O = 'O'
    AB = 'AB'

列挙型クラスの方が、if文で書かれるよりも 血液型はどんな値が入っているのか 明示的に表現できて個人的には好きです。

血液型やトランプの記号(ハートとかダイヤとか)など

どの値を使うのか決まっているものを列挙型で表現するという選択肢を知っていると

安全でわかりやすいクラスを書くのに役に立ちます。


enumモジュールの基本的な使い方

import enum


class BloodType(enum.Enum):
       A = 'A型'
       B = 'B型'
       O = 'O型'
       AB = 'AB型'
>>> import app.example
>>>
>>>
>>> app.example.BloodType.O
<BloodType.O: 'O型'>
>>> app.example.BloodType.O.value
'O型'
>>> app.example.BloodType.O.name
'O'


ダンダーメソッドや一般的なメソッドもクラスに作れる

例えばこんな風に。

class BloodType(enum.Enum):
    A = 'A型'
    B = 'B型'
    O = 'O型'
    AB = 'AB型'

    def __str__(self):
        return self.name

    @classmethod
    def show_all(cls) -> List[str]:
        return list(map(lambda c: c.value, cls))


enum.auto, enum.unique, enum.IntEnum, enum.Enumの使いどころ

enum.auto

先程の血液型の例ですが、AかBかOかABかそれぞれ識別できれば良いだけで値は何でも良いとします。

そのような時に、enum.autoを使うと値をPythonが決定してくれます。

import enum


class BloodType(enum.Enum):
    A = enum.auto()
    B = enum.auto()
    O = enum.auto()
    AB = enum.auto()

自動で、1, 2, 3, 4...と振ってくれます。


enum.unique

通常のEnumだと、列挙型メンバーに同じ名前をつけることは出来ませんが、同じ値に設定することはできます。

>>> import enum
>>>
>>>
>>> class BloodType(enum.Enum):
...     A = 'A'
...     A = 'B'
...     
... 
Traceback (most recent call last):
  File "<input>", line 1, in <module>
    class BloodType(enum.Enum):
  File "<input>", line 3, in BloodType
    A = 'B'
  File "/home/mizzsugar/.pyenv/versions/3.8.0/lib/python3.8/enum.py", line 95, in __setitem__
    raise TypeError('Attempted to reuse key: %r' % key)
TypeError: Attempted to reuse key: 'A'


>>> import enum
>>>
>>>
>>> class BloodType(enum.Enum):
...     A = 'A'
...     A = 'B'

enum.uniqueは、valueが重複してほしくない時に使います。

>>> import enum
>>>
>>>
>>> @enum.unique
... class BloodType(enum.Enum):
...     A = 'A'
...     B = 'A'
...     
... 
Traceback (most recent call last):
  File "<input>", line 2, in <module>
    class BloodType(enum.Enum):
  File "/home/mizzsugar/.pyenv/versions/3.8.0/lib/python3.8/enum.py", line 860, in unique
    raise ValueError('duplicate values found in %r: %s' %
ValueError: duplicate values found in <enum 'BloodType'>: B -> A


enum.IntEnum

enum.Enumではvalueがintやfloatなどの数値であっても、int型やfloat型のオブジェクトと比較ができません。

>>> import enum
>>> class BloodType(enum.Enum):
...     A = 1
...     B = 2
...     O = 3
...     AB = 4
...
...
>>> BloodType.A == 1
False

しかし、IntEnumだと比較することができます。

>>> import enum
>>> class BloodType(enum.IntEnum):
...     A = 1
...     B = 2
...     O = 3
...     AB = 4
...
...
>>> BloodType.A == 1
True

例えば、他の場所から取得した数値とIntEnumの値を比較または計算したい時に役に立ちます。(ex: データベースから数値を取得して比較したい時など)


enumモジュールを使ったリファクタリング

「現場で役立つシステム設計の原則」 の第2章で書かれている申請管理システムをPythonで書き直した例を使って締めたいと思います。

申請管理システムの、申請状況を管理する部分を書きます。

ざっくり説明すると、最初になにかしたらの申請があって審査が始まったらその申請のステータスは「審査中」になる、みたいなやつです。

ここでのポイントは、ステータスが「承認済み」の申請が「差し戻し中」に変わらないように、といった意図しないステータスに遷移しないようにもしないといけないところです。

申請管理を適切な英語に直す時間がなかったので日本語プログラミングです。

https://www.amazon.co.jp/dp/B073GSDBGT/ref=dp-kindle-redirect?_encoding=UTF8&btkr=1


まず、Enum書かないで自分が考えた中で一番読みにくい例で。

def 遷移可能かチェックする(_from: str, to: str) -> bool:
    if _from == '審査中' and to in {'差し戻し中', '承認済み'}:
        return True
    if _from == '差し戻し中' and to in {'審査中', '終了'}:
        return True
    if _from == '承認済み' and to in {'実施中', '終了'}:
        return True
    if _from == '実施中' and to in {'終了', '中断中'}:
        return True
    if _from == '中断中' and to in {'実施中', '終了'}:
        return True
    return False

これだと、if文が煩雑で読みにくいです。


フェアじゃないのでちょっと変えます。

def 遷移可能かチェックする(_from: str, to: str) -> bool:
    if _from == '審査中' and to in {'差し戻し中', '承認済み'}:
        return True
    if _from == '差し戻し中' and to in {'審査中', '終了'}:
        return True
    if _from == '承認済み' and to in {'実施中', '終了'}:
        return True
    if _from == '実施中' and to in {'終了', '中断中'}:
        return True
    if _from == '中断中' and to in {'実施中', '終了'}:
        return True
    return False


先程のよりはスッキリしました。

ただ、この例だとどのような状態が存在するのかを管理するところがありません。

文字列で制御しているため、関係のない文字列が入力される可能性があります。

また、遷移可能かチェックする関数の他にも申請ステータスを使って計算や判定をする場合

文字列だと自由に書けてしまうため

関数1で書かれているステータスの種類と関数2で書かれているステータスの種類が違う…なんてことが起こりえます。


そこで、どのような種類のステータスがあるのかは、Enumクラスで管理することにします。

import dataclasses
import enum
from typing import (
    Dict,
    Set,
)


class 申請ステータス(enum.Enum):
    審査中 = enum.auto()
    差し戻し中 = enum.auto()
    承認済み = enum.auto()
    実施中 = enum.auto()
    中断中 = enum.auto()
    終了 = enum.auto()


@dataclasses.dataclass
class 申請ステータス遷移:
    allowed: Dict[申請ステータス, Set[申請ステータス]]

    def __init__(self) -> None:
        # 本当はFinalつけて再代入負不可にしたかったけど
        # Cannot redefine an existing name as final と言われてできませんでした。
        # init=Falseにしてもできませんでした。
        # allowed: Dict[申請ステータス, Set[申請ステータス]] = { ... } は
        # Dictがmutableであるため
        # mutable default <class 'dict'> for field allowed is not allowed となります。
        # PythonにImmutalbeなMappingがあれば良いのですがないようです。
        self.allowed = {
            申請ステータス.審査中: {
                申請ステータス.差し戻し中, 申請ステータス.承認済み
            },
            申請ステータス.差し戻し中: {
                申請ステータス.審査中, 申請ステータス.終了
            },
            申請ステータス.承認済み: {
                申請ステータス.実施中, 申請ステータス.終了
            },
            申請ステータス.実施中: {
                申請ステータス.終了, 申請ステータス.中断中
            },
            申請ステータス.中断中: {
                申請ステータス.実施中, 申請ステータス.終了
            },
        }

    def can_transit(
        self,
        _from: 申請ステータス,
        to: 申請ステータス
    ) -> bool:
        allowd_states: Final = self.allowed[_from]
        return to in allowd_states


これだと、どのような申請ステータスがあるのかをEnumクラスで一元管理しているので

申請ステータスを利用する時はこのクラスを使うようになります。

申請ステータスの種類が増えたり変更したら

申請ステータスのEnumクラスを変更すれば良いです。

文字列だと申請ステータスを利用している関数一つ一つを確認して変更する必要がありますが、それがなくなるので変更がより安全になります。



以上、Pythonenumモジュールの紹介でした。

使ったことがない方も、便利なのでぜひ使っていただけたら、と思います!