在 Python Enum 中使用 ClassVar

研究了一下文件， 沒看到能剛好 fit 手邊需求的用法＠＠...

以官方文件的示範為例，建立一個 Enum:

class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3


print(Color._member_map_)
# {'RED': <Color.RED: 1>, 'GREEN': <Color.GREEN: 2>, 'BLUE': <Color.BLUE: 3>}

如果想要增加一個 ClassVar _favorite 紀錄某些 member, 例如：

class Color2(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3
    _favorite = {RED, BLUE}


print(Color2._member_map_)
# {'RED': <Color2.RED: 1>, 'GREEN': <Color2.GREEN: 2>, 'BLUE': <Color2.BLUE: 3>, '_favorite': <Color2._favorite: {1, 3}>}

但是 _favorite 會被當做 enum member ...

如果要避免， 可以使用官方文件的介紹的 _ignore_：

a list of names, either as a list or a str, that will not be transformed into members, and will be removed from the final class

class Color3(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3
    _ignore_ = ['_favorite']
    _favorite = {RED, BLUE}


print(Color3._member_map_)
# {'RED': <Color3.RED: 1>, 'GREEN': <Color3.GREEN: 2>, 'BLUE': <Color3.BLUE: 3>}

ok! 看起來沒問題了

拿來應用吧！

class Color4(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3
    _ignore_ = ['_favorite']
    _favorite = {RED, BLUE}

    @property
    def is_favorite(self):
        return self.value in self._favorite


print(Color4.RED.is_favorite)
# AttributeError: 'Color4' object has no attribute '_favorite'

...... ?

---

翻一下 source code 發現， _ignore_ 相關的處理， 是在 Enum 的 metaclass EnumType 的 __new__ 之中：

# python3.*/enum.py
# ...
class EnumType(type):
    # ...
    def __new__(metacls, cls, bases, classdict, ...):
        # ...
        classdict.setdefault('_ignore_', []).append('_ignore_')
        ignore = classdict['_ignore_']
        for key in ignore:
            classdict.pop(key, None)
        # ...
    # ...
# ...
class Enum(metaclass=EnumType):
    # ...
# ...

ok, 看起來應該沒有辦法使用 _ignore_ 了

---

剛好前面噴的是 AttributeError ，一副就是在暗示用 descriptor 是可行的

試試看 descriptor:

class Descriptor:
    def __init__(self, value):
        self.value = value

    def __get__(self, obj, objtype=None):
        return self.value


class Color5(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3
    _favorite = Descriptor({RED, BLUE})

    @property
    def is_favorite(self):
        return self.value in self._favorite


print(Color5.RED.is_favorite)
# True
print(Color5.GREEN.is_favorite)
# False
print(Color5._member_map_)
# {'RED': <Color5.RED: 1>, 'GREEN': <Color5.GREEN: 2>, 'BLUE': <Color5.BLUE: 3>}

可以！

---

補 generic type, 做完整一點：

import typing as t
from enum import Enum


T = t.TypeVar('T')


class EnumClassVar(t.Generic[T]):
    def __init__(self, value: T):
        self.value = value

    def __get__(self, obj, objtype=None):
        return self.value


class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3
    _favorite = EnumClassVar({RED, BLUE})

    @property
    def is_favorite(self):
        return self.value in self._favorite

平常好像都沒實作 descriptor 這次不巧派上用場惹～