Atsushi2022の日記

データエンジニアリングに関連する記事を投稿してます

PythonでのEnum

標準ライブラリにあるEnumクラスにより、Enumを使用することができる。

列挙型は定数を表すために使われるため、Enumのメンバの名前は大文字にすることが推奨される。

次のコードで挙動を確認することができる。詳細については列挙型 HOWTOにわかりやすくまとめられている。

from enum import Enum, auto


class Colors(Enum):
    """
    >>> print(Colors.GREEN)
    Colors.GREEN
    >>> Colors.GREEN
    <Colors.GREEN: 2>
    >>> Colors(2)
    <Colors.GREEN: 2>
    >>> Colors["GREEN"]
    <Colors.GREEN: 2>
    >>> Colors(2).name
    'GREEN'
    >>> Colors(2).value
    2
    >>> Colors["MIDORI"] # MIDORIはGREENの別名
    <Colors.GREEN: 2>
    >>> Colors.MIDORI # MIDORIはGREENの別名
    <Colors.GREEN: 2>
    >>> [color for color in Colors]
    [<Colors.RED: 1>, <Colors.GREEN: 2>, <Colors.BLUE: 3>]
    >>> {color.name: color.value for color in Colors}
    {'RED': 1, 'GREEN': 2, 'BLUE': 3}
    """
    RED = 1
    GREEN = 2
    BLUE = 3
    MIDORI = 2


class Weekday(Enum):
    """
    >>> Weekday.FIRST_DAY_OF_WEEK
    <Weekday.MONDAY: 1>
    >>> from datetime import datetime
    >>> day = datetime.strptime("2024-03-23", '%Y-%m-%d')
    >>> Weekday.from_date(day)
    <Weekday.SATURDAY: 6>
    """
    MONDAY = 1
    TUESDAY = 2
    WEDNESDAY = 3
    THURSDAY = 4
    FRIDAY = 5
    SATURDAY = 6
    SUNDAY = 7
    FIRST_DAY_OF_WEEK = 1

    @classmethod
    def from_date(cls, date):
        return cls(date.isoweekday())


class Vehicles(Enum):
    """
    >>> [member.value for member in Vehicles]
    [1, 2, 3]
    """
    Car = auto()
    Airplane = auto()
    Bicycle = auto()


class AutoName(Enum):
    @staticmethod
    def _generate_next_value_(name, start, count, last_values):
        return name


class Ordinal(AutoName):
    """
    >>> [member.value for member in Ordinal] # _generate_next_value_()でauto()で返す値をオーバーライド
    ['NORTH', 'SOUTH', 'EAST', 'WEST']
    """
    NORTH = auto()
    SOUTH = auto()
    EAST = auto()
    WEST = auto()


if __name__ == "__main__":
    import doctest

    doctest.testmod()

参考

列挙型 HOWTO

列挙型のサポート