有一个 Python 泛型类型注释的问题

from enum import Enum as _Enum
from typing import TypeVar, Type, Generic, Callable


class BaseModel(_BaseModel):
    pass


T = TypeVar('T')
GT = TypeVar('GT')


# 扩展 enum
class EnumBase(Generic[GT], _Enum):
    # 额外对象
    extra_obj: GT

    def __new__(
            cls: Type[T], value: str, extra_obj: GT = None,
    ) -> T:
        obj = object.__new__(cls)
        obj._value_ = value
        obj.extra_obj = extra_obj
        return obj


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


# 这里应该如何注释泛型? 这样写 pycharm 是会提示类型的, 但是会报错
class EnumTest(EnumBase[B]):
    A = 'A', B(1)
    B = 'B', B(2)


foo = EnumTest.A.extra_obj
print(foo.value)

报错内容是

  File "xxx", line 35, in <module>
    class EnumTest(EnumBase[B]):
  File "xxx\enum.py", line 408, in __getitem__
    return cls._member_map_[name]
KeyError: <class '__main__.B'>

也就是说 EnumBase[B]会直接触发 EnumMeta 的__getitem__, 不知道有没有什么好的解决办法

第 1 条附言 · 2023-04-13 23:29:00 +08:00

解决了

10楼的建议不错, 我应该将那个能附加额外对象的类变成一个Mixin类，所以代码用例修改成这样

from enum import Enum

from typing import TypeVar, Type, Generic, NamedTuple

T = TypeVar('T')
GT = TypeVar('GT')


class ExtraEnumMixin(Generic[GT]):
    extra_obj: GT

    def __new__(
            cls: Type[T], value, extra_obj: GT = None,
    ) -> T:
        obj = object.__new__(cls)
        obj._value_ = value
        obj.extra_obj = extra_obj
        return obj


class Name(NamedTuple):
    en_name: str
    zh_name: str


class EnumTest(ExtraEnumMixin[Name], Enum):
    A = 'A', Name('apple', '苹果')
    B = 'B', Name('banana', '香蕉')


na = EnumTest.A.extra_obj
print(na.en_name, na.zh_name)

13 条回复 • 2023-04-15 23:59:23 +08:00

antness

2023-04-10 00:41:25 +08:00

```python
class EnumTest(EnumBase):
A: 'EnumBase[B]' = 'A', B(1)
B: 'EnumBase[B]' = 'B', B(2)
```
这样可以实现你的需求吗

antness

2023-04-10 00:45:04 +08:00

或者
```python
class EnumTest(EnumBase):
A = 'A', B(1) # type: EnumTest[B]
B = 'B', B(2) # type: EnumTest[B]
```

NoOneNoBody

2023-04-10 00:55:22 +08:00

typing.ClassVar

egen

2023-04-10 01:04:40 +08:00

感觉这种问题很适合用 gpt 来做分析，用 gpt 几个来回分析后给出了可行的方案

以下是 gpt 的回复：
根据你的代码以及产生的错误，我为您分析了原因并提供了一种解决方案。问题的根源在于 EnumMeta （_Enum 基类的元类）没有为泛型类提供支持。因此，我们可以通过创建一个自定义的元类来解决这个问题。

这里是修改后的代码：
```
from enum import Enum as _Enum, EnumMeta as _EnumMeta
from typing import TypeVar, Type, Generic, Any

T = TypeVar('T')
GT = TypeVar('GT', bound=Type)

class _CustomEnumMeta(_EnumMeta):
def __getitem__(self, item: Any) -> Any:
if not isinstance(item, str): # 当处理泛型时，我们可以直接返回自身
return self
return super().__getitem__(item)

class EnumBase(Generic[GT], _Enum, metaclass=_CustomEnumMeta):
extra_obj: GT

def __new__(
cls: Type[T], value: str, extra_obj: GT = None,
) -> T:
obj = object.__new__(cls)
obj._value_ = value
obj.extra_obj = extra_obj
return obj

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

class EnumTest(EnumBase[B]):
A = 'A', B(1)
B = 'B', B(2)

foo = EnumTest.A.extra_obj
print(foo.value)
```

就是成本高了些，几个回合对话后总成本大概：¥3

anaf

2023-04-10 02:47:50 +08:00

啊这对于从 2.7 开始 3.7 结束的人来说有些看不懂这个代码了

ruanimal

2023-04-10 09:59:56 +08:00

@anaf 确实，而且在动态语言里这么写有点邪道了

thinszx

2023-04-10 13:40:45 +08:00

啊...第一次知道原来 python 也有泛型...

gujigujij

2023-04-11 00:02:02 +08:00

好像不行. 另外你的扩展 Enum 看着有点奇怪, 如果你想有额外对象, 可以这么写:

```python
class B:
def __init__(self, value):
self.value = value

class EnumBase( _Enum):
def __init__(self, code: str, extra_obj: B = None):
self.code: str = code
self.extra_obj: B = extra_obj

class EnumTest(EnumBase):
A = 'A', B(1)
B = 'B', B(2)

res = EnumTest.A.extra_obj.value
print(res)

```

gujigujij

2023-04-11 00:03:26 +08:00

```# markdown test
print("hello world")```

whitewinds

2023-04-12 00:54:46 +08:00

# 扩展 enum
class EnumBase(Generic[GT]): # <---
# 额外对象
extra_obj: GT

def __new__(cls: Type[T], value: str, extra_obj: GT = None,) -> T:
obj = object.__new__(cls) # type: Any
obj._value_ = value
obj.extra_obj = extra_obj
return obj

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

class EnumTest(EnumBase[B], _Enum): # <---
A = 'A', B(1)
B = 'B', B(2)

foo = EnumTest.A.extra_obj
print(foo.value)

bestcondition

2023-04-13 22:54:51 +08:00

@whitewinds 这个 mixin 写的好！受教了！

bestcondition

2023-04-13 23:03:32 +08:00

@egen 可以用，gpt 还真是挺厉害的！不过完全为了泛型注解而改元类的 getitem 感觉有点不妥

egen

2023-04-15 23:59:23 +08:00

@bestcondition gpt 只会傻傻的让代码可以通过，但是不一定能用正确的方法解决，用来解决一些问题还勉强可以，用来学习可能会被带歪，哈哈