python - 我可以将 Python 枚举编译成 numba jitclass 吗？

Question

我正在尝试寻找如何将@jitclass 与枚举类一起使用。引用手册明确表示它们是 supported但我无法弄清楚，而且我在任何地方都找不到代码示例。

当我尝试运行时

from numba import jitclass
from numba.types import string

from enum import Enum

type_spec = [
    ('A', string),
    ('B', string)
]


@jitclass(type_spec)
class Type(Enum):
    A = 'A'
    B = 'B'

我得到 TypeError: class members are not yet supported: _missing_, name, value, _convert, _member_names_, _member_map_, _member_type_, _value2member_map_, A, B

当我只是尝试在没有规范的情况下进行编译时，如下所示:

from numba import jitclass

from enum import Enum


@jitclass
class Type(Enum):
    A = 'A'
    B = 'B'

我得到一个 AttributeError: items当我调用构造函数时。我似乎无法找到将我的 Enum 类转换为 jitclass 的正确语法。我怎样才能做到这一点？

Answer 1

您可以在 njit 函数中使用 Enum/IntEnum 而无需在枚举类上使用 jitclass:

import numba
from enum import Enum, IntEnum

class MyEnum(Enum):
    Red = 123
    Blue = 456

class IEnum(IntEnum):
    one = 1
    two = 2

@numba.njit
def check_isblue(e):
    assert e == MyEnum.Blue

@numba.njit
def use_cases(a):
    """The following use-cases are supported"""
    # comparison
    a == MyEnum.Red
    a != MyEnum.Red
    a is MyEnum.Red
    a is not MyEnum.Red
    # getitem
    a is MyEnum["Red"]
    # conditional
    MyEnum.Red if True else MyEnum.Blue
    # coersion to int
    IEnum.one + 1


check_isblue(MyEnum.Blue)
use_cases(MyEnum.Red)

其他用例，例如与 IntEnum 值进行矢量化比较(也支持通过 @vectorize 装饰器)。

这是在 numba/tests/test_enums.py 中包含的 pull-request 1829 测试模块中“记录”的。