何かを書き留める何か

数学や読んだ本について書く何かです。最近は社会人として生き残りの術を学ぶ日々です。

Pythonの列挙型と型チェッカについて

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

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の本格的な入門書や専門書を監訳したり翻訳したりしました。 是非とも書店で手に取ったり、直販サイトで購入していただきたい。

www.oreilly.co.jp

www.oreilly.co.jp

www.oreilly.co.jp

『データエンジニアリングの基礎』の査読を担当しました

Reviewer

データエンジニアリングとはこれ

2024年3月27日にオライリージャパンから『Fundamentals of Data Engineering』の邦訳である『データエンジニアリングの基礎』が発売される。

www.oreilly.co.jp

本書は「データエンジニアリングライフサイクル」という概念を軸にデータシステムの要件を整理、システムの構築、データエンジニアの立ち位置などを詳細に論じたものである。 データサイエンティストではなくデータ分析の基盤を作るデータエンジニアの本ではあるものの、普段はWebシステムを作っているのでそこまでデータサイエンス関係ないんだよな、という人でも面白く読めるはずである。 特定のツールに依存しない説明であり、結果的にRDSしか使わないという選択をしたとしても本来はここまで考えないといけないのか、という記述の連続である。

データエンジニアに限らず、何らかのソフトウェアシステムに携わる人ならば是非とも読んでみてほしい。

Pythonの io.BytesIO と zipfile.ZipFile の組み合わせ

Python

with文がやってくれるのはどこまでなのか

インターネット経由で取得したZIPファイルを手元で加工する、という状況を考える。 たとえば、次のようなコードを書いたとする。 URL は ZIPファイルを取得できるものならば何でもよいが、今回は環境に配慮してローカルに1回だけダウンロードしてhttp.serverで簡易Webサーバを立てることでキャッシュしている。

import io
import urllib.request
import zipfile

URL = "http://localhost:8000/python-3.12.2-embed-amd64.zip"

with urllib.request.urlopen(URL) as f:
    content: bytes = f.read()

with zipfile.ZipFile(io.BytesIO(content)) as zf:
    names: list[str] = zf.namelist()

for name in names:
    print(name)

urllib.request.urlopen() 経由で取得したものはbytesであり、ファイルオブジェクトではない。 zipfile.ZipFile の第一引数file はファイルのパス、ファイルオブジェクト、 pathlib.Path のいずれかである必要があるため、urllib.request.urlopen() 経由で取得したものをそのまま渡すわけにはいかない。 そのため、 io.Bytes() を経由することでインメモリーストリームとして渡す。

このエントリの主題は、上記のコードにおける io.BytesIO(content) の取り扱いについてである。 with zipfile.ZipFile(io.BytesIO(content)) as zf: と記述しているので、 names: list[str] = zf.namelist() が完了したら zipfile.ZipFile(io.BytesIO(content)) の close() メソッドが呼ばれる。 その際に、中身である io.BytesIO(content) は閉じられるのだろうか。

公式ドキュメントとソースコードを巡る冒険

公式ドキュメントには

The buffer is discarded when the close() method is called.

とある通り、 BytesIO の close() メソッドが呼ばれなければインメモリーストリームのバッファは解放されない。

まず、 ソースコード にある __init__() から説明に必要な個所を抜き出す。

class ZipFile:
    fp = None                   # Set here since __del__ checks it

    def __init__(self, file, mode="r", compression=ZIP_STORED, allowZip64=True,
                 compresslevel=None, *, strict_timestamps=True, metadata_encoding=None):
        """Open the ZIP file with mode read 'r', write 'w', exclusive create 'x',
        or append 'a'."""
        ...

        # Check if we were passed a file-like object
        if isinstance(file, os.PathLike):
            file = os.fspath(file)
        if isinstance(file, str):
            # No, it's a filename
           ...
        else:
            self._filePassed = 1
            self.fp = file
            self.filename = getattr(file, 'name', None)
        self._fileRefCnt = 1
        ...

ここで認識してほしいのは以下の3点である。

  • zipfile.ZipFile の第一引数にファイルオブジェクトを渡すと self._filePassed = 1 となる。この変数は __init__()でのみ定義され、変更されない。
  • zipfile.ZipFile の第一引数にファイルオブジェクトを渡すと self.fp = file 、つまり self.fp にBytesIOインスタンスが入る。
  • zipfile.ZipFile の第一引数にファイルオブジェクトを渡すと self._fileRefCnt = 1 となる。今回のスニペットだと self._fileRefCnt = 1 は1のままである。

次に ソースコード にある close() メソッドを確認する。

    def close(self):
        """Close the file, and for mode 'w', 'x' and 'a' write the ending
        records."""
        if self.fp is None:
            return
        try:
            ...
        finally:
            fp = self.fp
            self.fp = None
            self._fpclose(fp)

ここでわかることは以下の点である。

  • 内部変数fpにself.fp 、つまり BytesIOインスタンスが入る。
  • その後、 self.fp は None となる。
  • self._fpclose(fp) が実行される。

そして、 ソースコード にある _fpclose()メソッドを確認する。

    def _fpclose(self, fp):
        assert self._fileRefCnt > 0
        self._fileRefCnt -= 1
        if not self._fileRefCnt and not self._filePassed:
            fp.close()
  • self._fileRefCnt = 1 なので、if 文に到達する際に `self._fileRefCnt == 0 となり、条件式の左辺は True である。
  • self._filePassed = 1 なので、条件式の右辺は False である。
  • よって、 if not self._fileRefCnt and not self._filePassed: は False となり、 fp.close() は実行されない。

以上より、「 io.BytesIO(content) は閉じられるのだろうか」という疑問は「閉じられません」となる。

では、閉じるにはどうすればよいのか、自分で明示的に閉じるしかない。

with io.BytesIO(content) as bs:
    with zipfile.ZipFile(bs) as zf:
        names: list[str] = zf.namelist()

こうすれば、 names: list[str] = zf.namelist() が実行された後、zipfile.ZipFileが閉じられ、 io.BytesIO が閉じられる。 これで安心して眠りにつくことができる。

ガベージコレクションの存在

しかし、このスニペット程度ならば、ガベージコレクション先生が気を利かしてくれるはずだ。

import io
import gc
import urllib.request
import zipfile

gc.enable()

URL = "http://localhost:8000/python-3.12.2-embed-amd64.zip"

with urllib.request.urlopen(URL) as f:
    content: bytes = f.read()

data = io.BytesIO(content)
with zipfile.ZipFile(data) as zf:
    names: list[str] = zf.namelist()

for name in names:
    print(name)

print(gc.is_tracked(data))

...
True

当然、 io.BytesIO(content) はガベージコレクタによって捕捉されており、必要に応じて回収されるだろう。

今後の課題

今回のスニペット程度ではガベージコレクタによって回収されるので、閉じ忘れても特に支障はないと思われる。 「ZIPファイルを手元で加工する」の内容如何では、ガベージコレクションによって回収されない状況が起き得るか、が自分の課題である。 また、 gc モジュールの使い方がイマイチわかっていないので、それの調査も課題である。