Repository: numba/numba

Issue:
Type inference fails when returning IntEnum
Returning an `IntEnum` from either a `vectorize` or `jit` wrapped function results in an a failure during type inference.

```python
# Saved as test.py
from enum import IntEnum

import numpy as np
import numba as nb


categories = ['NA', 'red', 'orange', 'yellow', 'blue', 'green', 'purple']
COLORS = IntEnum('Colors', categories)


def complement(a):
    if a == COLORS.red:
        return COLORS.green
    elif a == COLORS.orange:
        return COLORS.blue
    elif a == COLORS.yellow:
        return COLORS.purple
    elif a == COLORS.green:
        return COLORS.red
    elif a == COLORS.blue:
        return COLORS.orange
    elif a == COLORS.purple:
        return COLORS.yellow
    return COLORS.NA


data = np.array([COLORS.red, COLORS.blue, COLORS.green], dtype='i8')


cases = [('Jit with no signature', nb.jit(nopython=True), 1),
         ('Jit with signature', nb.jit('i8(i8)', nopython=True), 1),
         ('Vectorize with no signature', nb.vectorize(nopython=True), data),
         ('Vectorize with signature', nb.vectorize('i8(i8)', nopython=True), data)]

for name, dec, arg in cases:
    f = dec(complement)
    try:
        f(arg)
    except:
        print(name + ": FAILED")
    else:
        print(name + ": PASSED")
```

```bash
(numba) jcrist numba $ python test.py
Jit with no signature: FAILED
Jit with signature: PASSED
Vectorize with no signature: FAILED
Vectorize with signature: PASSED
```

Note that swapping out `COLORS.NA` for `1` causes all to pass.


Discussion:

Resolving patch:
diff --git a/numba/numpy_support.py b/numba/numpy_support.py
--- a/numba/numpy_support.py
+++ b/numba/numpy_support.py
@@ -130,6 +130,8 @@ def as_dtype(nbtype):
         return np.dtype('%s%d' % (letter, nbtype.count))
     if isinstance(nbtype, types.Record):
         return nbtype.dtype
+    if isinstance(nbtype, types.EnumMember):
+        return as_dtype(nbtype.dtype)
     raise NotImplementedError("%r cannot be represented as a Numpy dtype"
                               % (nbtype,))
 


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,8 +4,9 @@
 
 from __future__ import print_function
 
+import numpy as np
 import numba.unittest_support as unittest
-from numba import jit
+from numba import jit, vectorize
 
 from .support import TestCase
 from .enum_usecases import Color, Shape, Shake, Planet, RequestError
@@ -50,6 +51,13 @@ def int_coerce_usecase(x):
         return x + Shape.circle
 
 
+def vectorize_usecase(x):
+    if x != RequestError.not_found:
+        return RequestError['internal_error']
+    else:
+        return RequestError.dummy
+
+
 class BaseEnumTest(object):
 
     def test_compare(self):
@@ -127,6 +135,12 @@ def test_int_coerce(self):
         for arg in [300, 450, 550]:
             self.assertPreciseEqual(pyfunc(arg), cfunc(arg))
 
+    def test_vectorize(self):
+        cfunc = vectorize(nopython=True)(vectorize_usecase)
+        arg = np.array([2, 404, 500, 404])
+        sol = np.array([vectorize_usecase(i) for i in arg], dtype=arg.dtype)
+        self.assertPreciseEqual(sol, cfunc(arg))
+
 
 if __name__ == '__main__':
     unittest.main()