让枚举 (一如既往地，可以说是) 更 Pythonic

作者：Harry，2020-10-27

了解更多关于这张图片的信息

好吧，这实际上和软件架构没什么关系，但是

我讨厌 枚举！

最近在不得不处理它们的时候，我一次又一次地这样想。

为什么？

class BRAIN(Enum):
    SMALL = 'small'
    MEDIUM = 'medium'
    GALAXY = 'galaxy'

你可能会问，这有什么问题呢？好吧，如果你愿意指责我想要字符串化类型一切，但：这些枚举看起来像字符串，但它们不是！

assert BRAIN.SMALL == 'small'
# nope, <BRAIN.SMALL: 'small'> != 'small'

assert str(BRAIN.SMALL) == 'small'
# nope, 'BRAIN.SMALL' != 'small'

assert BRAIN.SMALL.value == 'small'
# finally, yes.

我猜有些人认为这是一个特性而不是一个 bug？但对我来说，这是一个无尽的烦恼之源。它们看起来像字符串！我将它们定义为字符串！为什么它们不像字符串一样工作呢！

只是一个常见的动机示例：通常你想用这些枚举做的事情是将它们转储到某个数据库列中。这种不完全是字符串的行为会导致你的 ORM 或 db-api 库疯狂地抱怨，并且在编写测试、自定义 SQL 等时会遇到无数的陷阱和令人挠头的问题。在这一点上，我真想把它们扔掉，直接使用普通的常量！

但是，Python 的 enum 模块的一个不错的承诺是 它是可迭代的。因此，不仅可以轻松地引用一个常量，还可以引用所有允许的常量的列表。也许这足以让人想要拯救它？

但是，再说一次，它并没有完全按照你想要的方式工作

assert list(BRAIN) == ['small', 'medium', 'galaxy']  # nope
assert [thing for thing in BRAIN] == ['small', 'medium', 'galaxy']  # nope
assert [thing.value for thing in BRAIN] == ['small', 'medium', 'galaxy']  # yes

这是一个 真的 让人无语的例子

assert random.choice(BRAIN) in ['small', 'medium', 'galaxy']
# Raises an Exception!!!

  File "/usr/local/lib/python3.9/random.py", line 346, in choice
    return seq[self._randbelow(len(seq))]
  File "/usr/local/lib/python3.9/enum.py", line 355, in __getitem__
    return cls._member_map_[name]
KeyError: 2

我不知道那里发生了什么。我们真正想要的是

assert random.choice(list(BRAIN)) in ['small', 'medium', 'galaxy']
# which is still not true, but at least it doesn't raise an exception

现在，如果你想将你的枚举进行鸭子类型化为整数，标准库确实提供了一个解决方案，IntEnum

class IBRAIN(IntEnum):
    SMALL = 1
    MEDIUM = 2
    GALAXY = 3

assert IBRAIN.SMALL == 1
assert int(IBRAIN.SMALL) == 1
assert IBRAIN.SMALL.value == 1
assert [thing for thing in IBRAIN] == [1, 2, 3]
assert list(IBRAIN) == [1, 2, 3]
assert [thing.value for thing in IBRAIN] == [1, 2, 3]
assert random.choice(IBRAIN) in [1, 2, 3]  # this still errors but:
assert random.choice(list(IBRAIN)) in [1, 2, 3]  # this is ok

这一切都很好，但我不想使用整数。我想使用字符串，因为这样当我查看我的数据库、打印输出或任何地方时，这些值都会有意义。

嗯，文档说你可以直接子类化 str 并创建你自己的 StringEnum，它会像 IntEnum 一样工作。但那是谎言

class BRAIN(str, Enum):
    SMALL = 'small'
    MEDIUM = 'medium'
    GALAXY = 'galaxy'

assert BRAIN.SMALL.value == 'small'  # ok, as before
assert BRAIN.SMALL == 'small'  # yep
assert list(BRAIN) == ['small', 'medium', 'galaxy']  # hooray!
assert [thing for thing in BRAIN] == ['small', 'medium', 'galaxy']  # hooray!
random.choice(BRAIN)  # this still errors but ok i'm getting over it.

# but:
assert str(BRAIN.SMALL) == 'small'   #NOO!O!O!  'BRAIN.SMALL' != 'small'
# so, while BRAIN.SMALL == 'small', str(BRAIN.SMALL)  != 'small' aaaargh

所以这就是我最终得到的

class BRAIN(str, Enum):
    SMALL = 'small'
    MEDIUM = 'medium'
    GALAXY = 'galaxy'

    def __str__(self) -> str:
        return str.__str__(self)
  • 这基本上避免了在你的代码中任何地方使用 .value 的需要
  • 枚举值以你期望的方式鸭子类型化为字符串
  • 你可以迭代 brain 并得到类似字符串的结果
  • 虽然 random.choice() 仍然有问题，但我将其留给读者作为练习
  • 并且类型提示仍然有效！
# both of these type check ok
foo = BRAIN.SMALL  # type: str
bar = BRAIN.SMALL  # type: BRAIN

示例代码 在一个 Gist 中，如果你想尝试一下。如果你发现任何更好的东西，请告诉我！