Repository: numba/numba

Issue:
Add support for `__getitem__` to `IntEnum`
Sometimes I have enums that aren't valid python identifiers. In these cases `IntEnum` supports indexing with a string:

```python
>>> from enum import IntEnum

>>> test = IntEnum('test', ['$$$', 'other'])
>>> test['$$$']
<test.$$$: 1>
```

However, this operator is currently not supported by numba.

```python
>>> import numba as nb
>>> @nb.jit('i8(i8)', nopython=True)
... def foo(x):
...     if x == test['$$$']:
...        return 10
...     return 5
---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)

<big traceback...>

TypingError: Failed at nopython (nopython frontend)
Invalid usage of getitem with parameters (IntEnum class<int64>(test), const('$$$'))
 * parameterized
```

Currently I have a rewrite rule that fixes this, but it'd be nice to add support for it in numba directly. I'm not sure how to go about doing this, but given a pointer I could maybe submit a PR myself :).


Discussion:

Resolving patch:
diff --git a/numba/targets/enumimpl.py b/numba/targets/enumimpl.py
--- a/numba/targets/enumimpl.py
+++ b/numba/targets/enumimpl.py
@@ -2,11 +2,10 @@
 Implementation of enums.
 """
 
-from llvmlite import ir
 
 from .imputils import (lower_builtin, lower_getattr, lower_getattr_generic,
                        lower_cast, lower_constant, impl_ret_untracked)
-from .. import typing, types, cgutils
+from .. import types
 
 
 @lower_builtin('==', types.EnumMember, types.EnumMember)
@@ -17,6 +16,7 @@ def enum_eq(context, builder, sig, args):
                                   (tu.dtype, tv.dtype), (u, v))
     return impl_ret_untracked(context, builder, sig.return_type, res)
 
+
 @lower_builtin('is', types.EnumMember, types.EnumMember)
 def enum_is(context, builder, sig, args):
     tu, tv = sig.args
@@ -30,7 +30,7 @@ def enum_is(context, builder, sig, args):
 
 
 @lower_builtin('!=', types.EnumMember, types.EnumMember)
-def enum_eq(context, builder, sig, args):
+def enum_ne(context, builder, sig, args):
     tu, tv = sig.args
     u, v = args
     res = context.generic_compare(builder, "!=",
@@ -42,6 +42,7 @@ def enum_eq(context, builder, sig, args):
 def enum_value(context, builder, ty, val):
     return val
 
+
 @lower_cast(types.IntEnumMember, types.Integer)
 def int_enum_to_int(context, builder, fromty, toty, val):
     """
@@ -49,6 +50,7 @@ def int_enum_to_int(context, builder, fromty, toty, val):
     """
     return context.cast(builder, val, fromty.dtype, toty)
 
+
 @lower_constant(types.EnumMember)
 def enum_constant(context, builder, ty, pyval):
     """
@@ -56,10 +58,22 @@ def enum_constant(context, builder, ty, pyval):
     """
     return context.get_constant_generic(builder, ty.dtype, pyval.value)
 
+
 @lower_getattr_generic(types.EnumClass)
-def enum_class_lookup(context, builder, ty, val, attr):
+def enum_class_getattr(context, builder, ty, val, attr):
     """
-    Return an enum member by name.
+    Return an enum member by attribute name.
     """
     member = getattr(ty.instance_class, attr)
     return context.get_constant_generic(builder, ty.dtype, member.value)
+
+
+@lower_builtin('static_getitem', types.EnumClass, types.Const)
+def enum_class_getitem(context, builder, sig, args):
+    """
+    Return an enum member by index name.
+    """
+    enum_cls_typ, idx = sig.args
+    member = enum_cls_typ.instance_class[idx.value]
+    return context.get_constant_generic(builder, enum_cls_typ.dtype,
+                                        member.value)
diff --git a/numba/typing/enumdecl.py b/numba/typing/enumdecl.py
--- a/numba/typing/enumdecl.py
+++ b/numba/typing/enumdecl.py
@@ -2,11 +2,9 @@
 Typing for enums.
 """
 
-import enum
-
-from numba import types, utils
+from numba import types
 from numba.typing.templates import (AbstractTemplate, AttributeTemplate,
-                                    signature, Registry, bound_function)
+                                    signature, Registry)
 
 registry = Registry()
 infer = registry.register
@@ -34,19 +32,32 @@ def generic_resolve(self, ty, attr):
             return ty.member_type
 
 
+@infer
+class EnumClassStaticGetItem(AbstractTemplate):
+    key = "static_getitem"
+
+    def generic(self, args, kws):
+        enum, idx = args
+        if (isinstance(enum, types.EnumClass)
+                and idx in enum.instance_class.__members__):
+            return enum.member_type
+
+
 class EnumCompare(AbstractTemplate):
 
     def generic(self, args, kws):
         [lhs, rhs] = args
         if (isinstance(lhs, types.EnumMember)
-            and isinstance(rhs, types.EnumMember)
-            and lhs == rhs):
+                and isinstance(rhs, types.EnumMember)
+                and lhs == rhs):
             return signature(types.boolean, lhs, rhs)
 
+
 @infer
 class EnumEq(EnumCompare):
     key = '=='
 
+
 @infer
 class EnumNe(EnumCompare):
     key = '!='


Test patch:
diff --git a/numba/tests/test_enums.py b/numba/tests/test_enums.py
--- a/numba/tests/test_enums.py
+++ b/numba/tests/test_enums.py
@@ -4,36 +4,44 @@
 
 from __future__ import print_function
 
-import enum
-
 import numba.unittest_support as unittest
 from numba import jit
 
-from .support import TestCase, tag
-from .enum_usecases import *
+from .support import TestCase
+from .enum_usecases import Color, Shape, Shake, Planet, RequestError
 
 
 def compare_usecase(a, b):
     return a == b, a != b, a is b, a is not b
 
-def global_usecase(a):
+
+def getattr_usecase(a):
     # Lookup of a enum member on its class
     return a is Color.red
 
+
+def getitem_usecase(a):
+    """Lookup enum member by string name"""
+    return a is Color['red']
+
+
 def identity_usecase(a, b, c):
     return (a is Shake.mint,
             b is Shape.circle,
             c is RequestError.internal_error,
             )
 
+
 def make_constant_usecase(const):
     def constant_usecase(a):
         return a is const
     return constant_usecase
 
+
 def return_usecase(a, b, pred):
     return a if pred else b
 
+
 def int_coerce_usecase(x):
     # Implicit coercion of intenums to ints
     if x > RequestError.internal_error:
@@ -70,7 +78,8 @@ def check_constant_usecase(self, pyfunc):
             self.assertPreciseEqual(pyfunc(arg), cfunc(arg))
 
     def test_constant(self):
-        self.check_constant_usecase(global_usecase)
+        self.check_constant_usecase(getattr_usecase)
+        self.check_constant_usecase(getitem_usecase)
         self.check_constant_usecase(make_constant_usecase(self.values[0]))