Python 的强类型枚举 (mypy)

Question

我使用

mypy

进行打字，但对于

Enum

效果不佳

我定义的大多数枚举只是整数映射

例如。

class Role(Enum):
    ENG = 0
    PRODUCT = 1
    BUSINESS = 2
    FOUNDER = 3

当我创建一个枚举并尝试获取其值（将其存储在

db

中）时，其类型被推断为 Any。

a = Role(2)
a.value  # inferred as Any instead of int

如何更改此设置以获得更好的打字效果？我不介意创建一个新的基类并使用它来代替 Enum。

Answer 1

您可以使用

IntEnum

和

StrEnum

来实现以下目的：

from enum import IntEnum

class Role(IntEnum):
    ENG = 0
    PRODUCT = 1
    BUSINESS = 2
    FOUNDER = 3

a = Role(2)
a.value  # inferred as int

例如：

import datetime
from enum import StrEnum


class Weekdays(StrEnum):
    sun = "Sunday"
    mon = "Monday"
    tue = "Tuesday"
    wed = "Wednesday"
    thu = "Thursday"
    fri = "Friday"
    sat = "Saturday"

    @classmethod
    def from_int(cls, value: int) -> "Weekdays":
        return cls(list(cls)[value])

    @classmethod
    def from_datetime(cls, dt: datetime.datetime | datetime.date) -> "Weekdays":
        return cls(dt.weekday()) # MyPy: Argument 1 to "Weekdays" has incompatible type "int"; expected "str"