如何在Numba jitclass的spec中声明Enum与自定义类？

在Numba v0.56.4中，jitclass的spec无法直接使用Python Enum类作为类型声明，需要通过以下方式处理：

核心思路 #

Numba对Enum的支持基于其底层数值类型，jitclass的成员变量需声明为Enum对应的数值类型（比如types.int64），同时在初始化或使用时手动处理Enum实例与数值的转换。

正确实现代码 #

from enum import Enum
from numba import types
from numba.experimental import jitclass

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

# spec中声明color为Enum底层的数值类型（int64）
spec = [('name', types.String), ('color', types.int64)]

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

    # 可选：添加方法将数值转回Enum（可在非jit代码中调用）
    def get_color_enum(self):
        return Color(self.color)

# 使用方式：传入Enum的value属性
red_paint = Paint("Red Paint", Color.RED.value)
# 若需要获取Enum实例
print(red_paint.get_color_enum())

为什么之前的尝试失败 #

• 直接在spec中写Color：jitclass的spec要求必须是Numba类型实例（如types.int64、types.String），而Color是Python的Enum类，不符合要求。
• 用types.int64但直接传Enum实例：Numba不会自动将Enum实例转换为对应的数值，必须手动提取.value属性传入。
• 用jitclass包装Enum：jitclass本身是Python类，不是Numba类型实例，spec中需要的是该jitclass的实例类型（如ColorWrapper.class_type.instance_type），但这种方式完全没必要——Numba已经原生支持在njit函数中使用Enum，无需额外包装。

额外优化 #

如果希望Paint的构造函数直接接受Enum实例，可以在jitclass外面套一个工厂函数：

def create_paint(name, color):
    return Paint(name, color.value)

# 直接传入Enum实例即可
blue_paint = create_paint("Blue Paint", Color.BLUE)

内容的提问来源于stack exchange，提问作者Yogesh