enum を拡張する extenum パッケージを作りました

背景

Python 3.4 から enum モジュールが追加されました。3.4 より低いバージョン (2.x も含む) 向けにバックポートされたパッケージが enum34 として PyPI に登録されています。

enum モジュールを使ったことがない方は、以下の記事が分かりやすいと思います。

• Python 3.4 から標準ライブラリに入る Enum 型が今からでも便利

CPython のコア開発者である Nick Coghlan 氏のノートに enum の標準ライブラリ化の背景や当時の議論の要約などが書かれています。

• Python, Enumerations and “Good Enough”

もともと標準になかったものなので Python の enum は Good Enough で満足と言えるかもしれません。

逆に足りないものは何だろう？Effective Java という、Java の世界のバイブルのような本があり、そこでは1章に渡って enum のベストプラクティスについて書かれています。

Java の世界で enum はコンパイル時の型安全性を保証したり、シングルトンを実装する方法であったりと応用範囲が広く、便利な機能なのでよく使われています。そこで Effective Java を眺めながら Python の enum にはない機能を拡張してみました。

extenum

標準の enum モジュールを拡張する extenum パッケージを作成しました。

インストール方法

$ pip install extenum

用語

公式ドキュメントの用語 によると、

3.4

>> from enum import Enum
>> class Color(Enum):

... red = 1
... green = 2
... blue = 3
...

> 注釈 用語
> * クラス `Color` は 列挙型 (または Enum) です。
> * 属性 `Color.red`、`Color.green`、その他は 列挙型のメンバー (または Enum メンバー) です。
> * 列挙型のメンバーは 名前 と 値 を持ちます (`Color.red` の名前は `red`、`Color.blue` の値は 3 など)。

列挙型で定義する定数を ***列挙型のメンバー (または Enum メンバー)*** と用語定義しています。他の言語では、列挙子や列挙定数と呼ばれたりもします。定数としての役割を期待しているので列挙定数の方が分かりやすい気もしますが、公式ドキュメントにあわせ、本稿では列挙型の定数を列挙型のメンバーと呼ぶようにします。

### 定数固定メソッド実装
Effective Java の名前が定数とありますが、この *定数* は列挙型のメンバーを指します。

extenum を作ろうと私が思ったきっかけがこの仕組みがないことでした。定数固定メソッドを探したときにそれっぽいものとして [asym-enum](https://github.com/AsymmetricVentures/asym-enum#methods-and-functions) を見つけましたが、

```py3:3.4
from asymm_enum.enum import Enum

class E(Enum):
    A = 1
    B = 2

    def get_label(self):
        ''' instance methods are attached to individual members '''
        return self.label

    @classmethod
    def get_b_label(cls):
        ''' class methods are attached to the enum '''
        return cls.B.label

これはちょっと格好悪いなと思って自分で作ることにしました。

extenum を使うと、以下のように実装できます。

constant_specific_enum.py

from extenum import ConstantSpecificEnum

class Operation(ConstantSpecificEnum):
    PLUS = '+'
    MINUS = '-'
    TIMES = '*'
    DIVIDE = '/'

    @overload(PLUS)
    def apply(self, x, y):
        return x + y

    @overload(MINUS)
    def apply(self, x, y):
        return x - y

    @overload(TIMES)
    def apply(self, x, y):
        return x * y

    @overload(DIVIDE)
    def apply(self, x, y):
        return x / y

きれいに実装できました。

Python にはオーバーロードの仕組みがないので @overload(CONSTANT) デコレーターで定数固定メソッドを実装します。

3.4

>>> from constant_specific_enum import Operation
>>> for name, const in Operation.__members__.items():
...     print(name, ':', const.apply(2, 4))
PLUS : 6
MINUS : -2
TIMES : 8
DIVIDE : 0.5

定数固定メソッドは、メソッドの名前により意図を伝えるのにも有効です。例えば、処理のステータス情報を保持する列挙型をみてみましょう。

3.4

from extenum import ConstantSpecificEnum

class Status(ConstantSpecificEnum):
    PREPARING = 1
    WAITING = 2
    RUNNING = 3

    @overload(PREPARING)
    def is_cancelable(self):
        return False

    @overload(WAITING)
    def is_cancelable(self):
        return True

    @overload(RUNNING)
    def is_cancelable(self):
        return True

それぞれの列挙型のメンバーがステータス値を指しますが、この処理のキャンセル処理を行う前にキャンセル可能なステータスかどうかを判別する必要があると仮定します。

定数固定メソッドがない場合、ある関数に渡された列挙型のメンバーが WAITING か RUNNING のどちらかであるかを調べています。

3.4

def cancel(status):
    if status in [Status.WAITING, Status.RUNNING]:
        # do cancel

これはキャンセル関数の実装者や読み手がステータス情報の列挙型の詳細を知っておく必要があり、ステータス情報が増えたときにこの関数内の条件を修正する必要があるかもしれません。

定数固定メソッドがある場合、このキャンセル関数の前処理のステータス判定はより自然にすっきりと書けます。

3.4

def cancel(status):
    if status.is_cancelable():
        # do cancel

前述したコードと比べると、キャンセル関数の実装者や読み手にも分かりやすく、ステータス情報が増えても列挙型の定数固定メソッドを追加するだけで済みます。

戦略 enum パターン

定数固定メソッド実装の応用例です。曜日と平日・休日の属性を列挙型をネストすることで、より保守性の高い実装を意図したものです。

strategy_enum_pattern.py

from extenum import ConstantSpecificEnum

class PayrollDay(ConstantSpecificEnum):

    class PayType(ConstantSpecificEnum):
        WEEKDAY = 1
        WEEKEND = 2

        @overload(WEEKDAY)
        def overtime_pay(self, hours, pay_rate):
            return 0 if hours <= 8 else (hours - 8) * pay_rate / 2

        @overload(WEEKEND)
        def overtime_pay(self, hours, pay_rate):
            return hours * pay_rate / 2

        def pay(self, hours_worked, pay_rate):
            base_pay = hours_worked * pay_rate
            overtime_pay = self.overtime_pay(hours_worked, pay_rate)
            return base_pay + overtime_pay

    MONDAY = PayType.WEEKDAY
    TUESDAY = PayType.WEEKDAY
    WEDNESDAY = PayType.WEEKDAY
    THURSDAY = PayType.WEEKDAY
    FRIDAY = PayType.WEEKDAY
    SATURDAY = PayType.WEEKEND
    SUNDAY = PayType.WEEKEND

    def pay(self, hours_worked, pay_rate):
        return self.value.pay(hours_worked, pay_rate)

Java だとプラクティスの1つとして便利ですが、Python だとやり過ぎ感が漂うような気もします (主観)。試しに実装してみましたが、良いとも悪いとも言えない微妙な例な気がします。

一応こんなこともできるんだ程度に見てください。

3.4

>>> from strategy_enum_pattern import PayrollDay
>>> PayrollDay.MONDAY.pay(10, 1000.0)
11000.0
>>> PayrollDay.WEDNESDAY.pay(8, 1000.0)
8000.0
>>> PayrollDay.SATURDAY.pay(10, 1000.0)
15000.0
>>> PayrollDay.SUNDAY.pay(8, 1000.0)
12000.0

暗黙の列挙型メンバー

列挙型のメンバーは、クラス変数の定義のため、以下のような列挙型の定義はできません。

3.4

class Color(Enum):
    red, green, blue

実行すると、NameError が発生する。

3.4

NameError: name 'red' is not defined

enum モジュールを追加した PEP 435 -Not having to specify values for enums- によると、

Cons: involves much magic in the implementation, which makes even the definition of such enums baffling when first seen. Besides, explicit is better than implicit.

暗黙的に値を設定するような定義ができない理由として、実装が魔術化する、最初に見たときに困惑する、暗黙よりも明示の禅に反するといったことが挙げられています。

これはこれで Python の文化としては正しいので別に構わないのですが、Nick Coghlan 氏のノートの Support for alternate declaration syntaxes に暗黙の列挙型メンバーを実装する方法について言及されています。

Implicit enums that don’t really support normal code execution in the class body, and allow the above to be simplified further. It’s another variant of the autonumbered example in the test suite, but one that diverges substantially from normal Python semantics: merely mentioning a name will create a new reference to that name. While there are a number of ways to get into trouble when doing this, the basic concept would be to modify __prepare__ on the metaclass to return a namespace that implements __missing__ as returning a custom sentinel value and overrides __getitem__ to treat repeating a name as an error:

Python の 黒魔術 メタプログラミングとしておもしろい方法なので実装してみました。

特殊メソッドの __prepare__ は Python 3 から追加されたもので、メタクラスが初期化されるときにそのクラスの名前空間を返す処理をフックします。__missing__ はディクショナリのサブクラスに対して、そのキーがディクショナリに存在しないとき KeyError のエラー処理をフックします。

先の __prepare__ で名前空間オブジェクトとして __missing__ をもつディクショナリを返せば、クラスの名前空間の定義をフックできるわけです。

何を言っているのか分からないのでコードにしてみましょう。

prepare_missing.py

class MyNameSpace(dict):
    def __missing__(self, key):
        self[key] = value = 'x'
        return value

class MyMetaclass(type):
    @classmethod
    def __prepare__(metacls, cls, bases):
        return MyNameSpace()

    def __new__(metacls, cls, bases, classdict):
        print('classdict is', classdict.__class__)
        return super().__new__(metacls, cls, bases, classdict)

このメタクラスを使ってクラスを定義します。

3.4

>>> from prepare_missing import MyMetaclass
>>> class Color(metaclass=MyMetaclass):
...     red, green, blue
... 
classdict is <class 'status.MyNameSpace'>
>>> Color.red
'x'
>>> Color.blue
'x'

クラスを定義するときにメタクラス MyMetaclass の __new__ メソッドが呼ばれ、クラスの名前空間 (classdict) に MyNameSpace が渡されています。MyNameSpace の __missing__ では、キーが存在しない、つまりクラスの名前空間に存在しない名前に対して、文字列の x を設定します。

この仕組みを応用すると ImplicitEnum は簡単に実装できます。

3.4

>>> from extenum import ImplicitEnum
>>> class Color(ImplicitEnum):
...     RED
...     GREEN
...     BLUE
...
>>> for name, const in Color.__members__.items():
...     print(name, ':', const.value)
...
RED : 1
GREEN : 2
BLUE : 3

余談ですが、前節の ConstantSpecificEnum の @overload デコレーターもこの仕組みを使って実装しています。このデコレーター定義そのものを隠してしまうこと自体は賛否両論あると思いますが、@classmethod や @staticmethod などがグローバルに扱えるように、隠してしまっても使う側から違和感はないかなぁと考えた次第です。

EnumSet

EnumSet は Enum 定数を含む集合 (set) を生成するユーティリティクラスのようなものです。Java の EnumSet を眺めながら実装してみました。

分かりやすい用途としてビット演算で実装していたフラグ操作のような処理を、EnumSet を使って実装すると人に易しい雰囲気になります。

3.4

>>> from enum import Enum
>>> class Mode(Enum):
...     READ = 4
...     WRITE = 2
...     EXECUTE = 1
... 

この Mode の列挙型を使って試してみましょう。

3.4

>>> from extenum import EnumSet

>>> EnumSet.all_of(Mode)  # Mode の全列挙型メンバーを含む EnumSet を生成
EnumSet({<Mode.READ: 4>, <Mode.WRITE: 2>, <Mode.EXECUTE: 1>})

>>> EnumSet.none_of(Mode)  # Mode の列挙型メンバーを受け取る空っぽの EnumSet を生成
EnumSet()

>>> enumset = EnumSet.of(Mode.READ, Mode.EXECUTE)  # 任意の列挙型メンバーを含む EnumSet を生成
>>> enumset
EnumSet({<Mode.READ: 4>, <Mode.EXECUTE: 1>})
>>> enumset.update(EnumSet.of(Mode.READ, Mode.WRITE))
>>> enumset
EnumSet({<Mode.READ: 4>, <Mode.WRITE: 2>, <Mode.EXECUTE: 1>})
>>> Mode.WRITE in enumset
True
>>> 2 in enumset
False

EnumSet にまとめてしまうことでモードやオプションなどの取り扱いが簡単になります。

まとめ

Java っぽい enum 拡張として extenum により、以下の機能を提供しました。

• 定数固定メソッド
• 暗黙の列挙型メンバー
• EnumSet

標準の Enum、EnumMeta、_EnumDict を継承して拡張しているため、おそらくは通常の列挙型としても素直に動くと思います。

他にもおもしろそうな enum のプラクティスがあれば教えてください。使うよりも作る方が目的だったりもします。