如何修复我的 Python 项目中类似 Enum 的类的 mypy“不兼容的返回值类型”？

Question

在我们的项目中，我们从不使用 Enum，我们使用 SimpleEnum，其定义如下：

class SimpleEnum:
    @classmethod
    def values(cls):
        return {
            getattr(cls, name)
            for name in dir(cls)
            if not name.startswith("_") and not callable(getattr(cls, name))
        }

现在是这样的：

class MySimpleEnum(SimpleEnum):
    completed = "completed"
    parse = "parse"


def foo() -> MySimpleEnum:
    return MySimpleEnum.parse

我收到 mypy 错误：

error: Incompatible return value type (got "str", expected "MySimpleEnum")  [return-value]

这是十万行项目中的超级简化用例。我不能只切换到 Enum。如何更改 SimpleEnum 以允许像我尝试做的那样进行代码提示（如果可能的话）？更改不得改变 SimpleEnum 的行为（因为它在项目中使用了数百次）。

Answer 1

这是一种潜在的笨拙解决方法：

import typing

class SimpleEnum:
    @classmethod
    def values(cls):
        return {
            getattr(cls, name)
            for name in dir(cls)
            if not name.startswith("_") and not callable(getattr(cls, name))
        }

MySimpleEnumValue = typing.NewType(str)

class MySimpleEnum(SimpleEnum):
    completed = MySimpleEnumValue("completed")
    parse = MySimpleEnumValue("parse")


def foo() -> MySimpleEnumValue:
    return MySimpleEnum.parse

这依赖于

typing.NewType

。新类型实际上并不是一个新类型，它只是用于静态分析目的，而

MySimpleEnumValue(x)

只是返回

x

，因此这些值实际上只是

str

的实例，但静态类型检查器应该考虑他们截然不同。

当然，它要求您为每个

typing.NewType

子类创建一个新的

SimpleEnum

。但这可能需要最少的重构。

但实际上，您应该重构以使用

enum.Enum

。不过，您可能只使用

enum.StrEnum

来使转换更容易，

import enum

class SimpleEnum(enum.StrEnum):
    @classmethod
    def values(cls):
        return set(cls.__members__.values())

注意，在这种情况下，成员是字符串的子类型：

In [19]: import enum
    ...:
    ...: class SimpleEnum(enum.StrEnum):
    ...:     @classmethod
    ...:     def values(cls):
    ...:         return set(cls.__members__.values())
    ...:

In [20]: class MySimpleEnum(SimpleEnum):
    ...:     completed = "completed"
    ...:     parse = "parse"
    ...:

In [21]: MySimpleEnum.values()
Out[21]: {<MySimpleEnum.completed: 'completed'>, <MySimpleEnum.parse: 'parse'>}

但它们是字符串：

In [22]: print(MySimpleEnum.completed)
completed

In [23]: MySimpleEnum.completed + " the job"
Out[23]: 'completed the job'

但是，如果您在代码中的任何位置测试

str

类型具体（标准库中有些地方会发生这种情况，请参阅文档中的注释），它就会中断。例如有事做：

if type(x) is str:
    ...

而不是

if isinstance(x, str):
   ...

另一个地方是，如果您在任何地方使用字符串的

repr

，这将使用枚举

__repr__

而不是

str.__repr__