Can I compile a Python enum into a numba jitclass?

user_8914113 0
Viewed 748

I am trying to find how to use @jitclass with an Enum class. The reference manual says explicitly that they are supported but I can't figure it out, and I cannot find a code example anywhere.

When I try running 

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'

I get TypeError: class members are not yet supported: _missing_, name, value, _convert, _member_names_, _member_map_, _member_type_, _value2member_map_, A, B

When I just try to compile without a spec, like below :

from numba import jitclass

from enum import Enum


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

I get an AttributeError: items when I call the constructor. I just can't seem to find the right syntax to turn my Enum class into a jitclass. How can I achieve this?

numba enums python
2 Answers

You can use Enum/IntEnum in an njit function without using jitclass on the enum class:

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)

Additional use-cases such as vectorized comparison with IntEnum values (via the @vectorize decorator are also supported).

This is "documented" in the numba/tests/test_enums.py test module that was included in pull-request 1829.

user_42563460

I am currently facing the same problem (not being able to jitclass a Python Enum). From what @user2357112-supports-Monica seems to imply, it seems that Numba documentation could be seriously improved.

"support" for enums was added in https://github.com/numba/numba/pull/1829, but that pull request was seriously lacking on the documentation side.

user_54385660
Related