Pythonの列挙型と型チェッカについて - 何かを書き留める何か

言いたいこと

Pythonの列挙型を定義する方法としてクラスと関数呼び出しの2種類存在する。
関数呼び出しで列挙型を定義するべき状況とは、列挙型を動的に制御したい状況である。
しかし、関数呼び出しで列挙型を定義すると、mypyを通して型チェックをする際に都合が悪い。
列挙型を動的に制御したい場合は列挙型を使わずに素直に辞書を使う。

列挙型とは

Pythonの公式ドキュメントには、列挙型について次のように説明している。

一意の値に紐付けられたシンボリックな名前の集合です

要は、まとめて扱いたい定数の集まりである。

列挙型の定義方法

Pythonには列挙型の定義方法として、クラスによる定義、関数呼び出しによる定義の2種類存在する。公式ドキュメントの例に基づいて、それぞれの例を示す。

クラスによる定義

import enum


class Planet(enum.Enum):
    MERCURY = (3.303e23, 2.4397e6)
    VENUS = (4.869e24, 6.0518e6)
    EARTH = (5.976e24, 6.37814e6)
    MARS = (6.421e23, 3.3972e6)
    JUPITER = (1.9e27, 7.1492e7)
    SATURN = (5.688e26, 6.0268e7)
    URANUS = (8.686e25, 2.5559e7)
    NEPTUNE = (1.024e26, 2.4746e7)

    def __init__(self, mass: float, radius: float) -> None:
        self.mass = mass  # in kilograms
        self.radius = radius  # in meters

    @property
    def surface_gravity(self) -> float:
        # universal gravitational constant  (m3 kg-1 s-2)

        G = 6.67300e-11
        return G * self.mass / (self.radius * self.radius)


print(Planet.EARTH.value)  # (5.976e+24, 6378140.0)
print(Planet.EARTH.surface_gravity)  # 9.802652743337129

関数呼び出しによる定義

import enum


class PlanetBase(enum.Enum):
    def __init__(self, mass: float, radius: float) -> None:
        self.mass = mass  # in kilograms
        self.radius = radius  # in meters

    @property
    def surface_gravity(self) -> float:
        # universal gravitational constant  (m3 kg-1 s-2)

        G = 6.67300e-11
        return G * self.mass / (self.radius * self.radius)


PLANET_VALUES: dict[str, tuple[float, float]] = {
    "MERCURY": (3.303e23, 2439700.0),
    "VENUS": (4.869e24, 6051800.0),
    "EARTH": (5.976e24, 6378140.0),
    "MARS": (6.421e23, 3397200.0),
    "JUPITER": (1.9e27, 71492000.0),
    "SATURN": (5.688e26, 60268000.0),
    "URANUS": (8.686e25, 25559000.0),
    "NEPTUNE": (1.024e26, 24746000.0),
}

Planet = PlanetBase("PLANETS", PLANET_VALUES)

print(Planet.EARTH.value)  # (5.976e+24, 6378140.0)
print(Planet.EARTH.surface_gravity)  # 9.802652743337129

関数呼び出しによる定義の使いどころ

上記の例は惑星に関する列挙型である。太陽系の惑星は水星から海王星までだが、人によっては冥王星を入れたいかもしれないし、G.Holstの組曲に合わせて地球を省きたいかもしれない。その場合は PLANET_VALUES を定義する際に if 文を入れるなどして、お好みの惑星を列挙できる。

一方、クラスの場合は別のクラスを定義するしかない。また、列挙型は継承（拡大）できない。実際にやってみるとわかる。

import enum


class E1(enum.Enum):
    M1 = enum.auto()
    M2 = enum.auto()


class E2(E1):
    M3 = enum.auto()
    
# TypeError: E2: cannot extend enumeration 'E1'

mypy（型チェッカ）と関数呼び出しによる定義の相性の悪さ

関数呼び出しによる定義は利点もあるが欠点もある。型チェッカとの相性が悪い。

たとえば、クラスによる定義では、適切に型ヒントを追加すればよい。

$ mypy --strict def_by_class.py
Success: no issues found in 1 source file

一方、関数呼び出しによる定義だと、引数を__call__()に渡しているはずなのに、__init__()の引数の型と一致しないというエラーが発生する。また、列挙型を生成した後も、存在するはずの属性EARTHが存在しないというエラーが発生する。

$ mypy --strict def_by_func.py
def_by_func.py:28: error: Argument 1 to "PlanetBase" has incompatible type "str"; expected "float"  [arg-type]
def_by_func.py:28: error: Argument 2 to "PlanetBase" has incompatible type "dict[str, tuple[float, float]]"; expected "float"  [arg-type]
def_by_func.py:30: error: "PlanetBase" has no attribute "EARTH"  [attr-defined]
def_by_func.py:31: error: "PlanetBase" has no attribute "EARTH"  [attr-defined]
Found 4 errors in 1 file (checked 1 source file)

今回の例は、列挙型のメンバの値を使う例なので少し複雑だが、単に enum.auto() を使ったとしても、属性が存在しないというエラーが発生する。そのため、相性の悪さとしては変わりない。

なお、これは mypy の問題ではない。 pyrightやpyreのPlaygroundで試してみたところ、やはり属性が存在しないというエラーが発生する。

github.com

代替手段としての辞書

『ロバストPython』の「8.2.2 列挙型を使うべきではない場合」にはこう書かれている。（なお、『ロバストPython』はこの記事の筆者が監訳しました。）

www.oreilly.co.jp

列挙型はコードの利用者に静的な選択肢を示すのに効果を発揮する。しかし、選択肢が実行時に決まるような場合は列挙型は使わないほうがよい。（中略）実行するたびに有効な値が変わるなら、コードの読み手はどの値が使えるかの判断が難しくなる。このような場合は、辞書を勧める。

ならば、実際に辞書で書いてみよう。

import dataclasses
import typing


PLANET_NAME = typing.Literal["MERCURY", "VENUS", "EARTH", "MARS", "JUPITER", "SATURN", "URANUS", "NEPTUNE"]


@dataclasses.dataclass
class PlanetValue:
    mass: float
    radius: float

    @property
    def surface_gravity(self) -> float:
        # universal gravitational constant  (m3 kg-1 s-2)

        G = 6.67300e-11
        return G * self.mass / (self.radius * self.radius)


PLANET_VALUES: dict[PLANET_NAME, PlanetValue] = {
    "MERCURY": PlanetValue(mass=3.303e23, radius=2439700.0),
    "VENUS": PlanetValue(mass=4.869e24, radius=6051800.0),
    "EARTH": PlanetValue(mass=5.976e24, radius=6378140.0),
    "MARS": PlanetValue(mass=6.421e23, radius=3397200.0),
    "JUPITER": PlanetValue(mass=1.9e27, radius=71492000.0),
    "SATURN": PlanetValue(mass=5.688e26, radius=60268000.0),
    "URANUS": PlanetValue(mass=8.686e25, radius=25559000.0),
    "NEPTUNE": PlanetValue(mass=1.024e26, radius=24746000.0),
}

print(PLANET_VALUES["EARTH"].mass)  # 5.976e+24
print(PLANET_VALUES["EARTH"].surface_gravity)  # 9.802652743337129

辞書だけでなく、データクラスも用いた。今回は typing.Literal を使ったが、外部から入力を受ける場合は何らかのバリデーション処理が必要になる。

これならば、必要に応じて惑星の構成要素を変更できるし、型チェッカも納得する。

$ mypy --strict def_by_dict.py
Success: no issues found in 1 source file

惑星の例では列挙型の方が優れているが、動的な列挙型を扱いたい場合は辞書の方が素直である。

言いたいこと（2回目）

Pythonの列挙型を定義する方法としてクラスと関数呼び出しの2種類存在する。
関数呼び出しで列挙型を定義するべき状況とは、列挙型を動的に制御したい状況である。
しかし、関数呼び出しで列挙型を定義すると、mypyを通して型チェックをする際に都合が悪い。
列挙型を動的に制御したい場合は列挙型を使わずに素直に辞書を使う。

宣伝

桜の季節にピッタリなPythonの本格的な入門書や専門書を監訳したり翻訳したりしました。是非とも書店で手に取ったり、直販サイトで購入していただきたい。