如何在 numba jitclass 规范中声明 Enum 和 Custom 类？

Question

我正在尝试 numba (v0.56.4)，我有一个 python 枚举 Color，我想在类 Paint 的 @jitclass 规范中声明它。请问可以告诉我正确的申报方式吗？

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

spec = [('name': types. String), ('color': Color)]

@jitclass(spec)
class Paint:
    def __init__(self, name, color):
        self.name = name
        self. Color = color

当我直接在规范中给出枚举颜色时，我收到错误 TypeError:

spec 值应该是 Numba 类型实例，得到

如果我尝试提供 int32 或 int64 而不是枚举 Color，则会收到错误：

('color', types.int64) # gives error 

无法将 Enum(Color) 转换为 int64。

如果我在 ColorWrapper 上使用 @jitclass 创建 Color 的包装类，我会收到错误

TypeError：规范值应该是 Numba 类型实例，得到

根据 numba 文档，我可以在 njit 函数中使用枚举，但我找不到使用 @jitclass 为另一个类的 @jitclass 规范声明枚举类型和自定义类的方法

谢谢

Answer 1

您可以将

Color

设为

enum.IntEnum

的子类（

Enum

类型，混合类型为

int

），以便可以将其转换为

int64

:

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

spec = [('name', string), ('color', int64)]