Numba jitclass 如何声明 Enum 与自定义类？

Numba 旨在通过即时 (JIT) 编译将 Python 代码转换为机器码，从而提高性能。@jitclass 装饰器允许用户定义可以被 Numba 编译的类，但正确声明类的属性类型至关重要。特别是在使用枚举 (Enum) 类型时，需要采用特定的方法才能使其与 Numba 兼容。

使用 enum.IntEnum 声明 Enum 类型

Numba 无法直接处理标准的 Python enum.Enum 类型。为了在 @jitclass 的 spec 中使用枚举，需要使用 enum.IntEnum。IntEnum 是 Enum 的一个子类，它继承了 int 类型，因此可以转换为 int64，这使得它与 Numba 兼容。

示例代码：

from enum import IntEnum
from numba import int64, string
from numba.experimental import jitclass

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

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

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

# 示例用法
paint = Paint("MyPaint", Color.RED)
print(paint.name)
print(paint.color)

代码解释：

1. from enum import IntEnum: 导入 IntEnum 类。
2. class Color(IntEnum):: 定义一个名为 Color 的枚举类，并继承自 IntEnum。
3. RED = 1, BLUE = 2, GREEN = 3: 定义枚举的成员及其对应的值。这些值必须是整数。
4. spec = [('name', string), ('color', int64)]: 在 spec 中，将 color 属性声明为 int64 类型。
5. @jitclass(spec): 使用 @jitclass 装饰器，并将 spec 传递给它。

注意事项

• 必须使用 IntEnum: 确保你的枚举类继承自 enum.IntEnum，而不是 enum.Enum。
• spec 中声明为 int64: 在 @jitclass 的 spec 中，将枚举类型的属性声明为 int64。
• 枚举值必须是整数: IntEnum 的成员的值必须是整数。

总结

通过继承 enum.IntEnum 并在 @jitclass 的 spec 中将枚举类型声明为 int64，可以有效地在 Numba 中使用枚举类型。这种方法允许 Numba 正确编译包含枚举类型的类，从而提高代码的性能。

虽然目前 Numba 还不支持直接声明自定义类在jitclass的spec中，但对于枚举类型，enum.IntEnum 提供了一个简单有效的解决方案。