Repository: numba/numba
Pull request: `IntEnumMember` support for  `np.empty`, `np.zeros`, and `np.ones`
Adds `IntEnumMember` support for `np.empty`, `np.zeros`, and `np.ones`.

Closes [issue 6631](https://github.com/numba/numba/issues/6631), as well as the implicitly-casted version of the code in question:
```python
import numba
from enum import IntEnum
import numpy as np
from numba import njit

class Matrix(IntEnum):
    r = 4
    c = 4

@njit
def test():
    return np.zeros((Matrix.r, Matrix.c), dtype='int64')

z = test()
print(z, z.dtype)
```

Is there a more general solution to this problem that would allow support for IntEnumMember in other places, or a good way to automate its inclusions throughout the codebase?

Problem statement:
Numpy does not support Enum / IntEnum
numba 0.52
numpy 1.19.5
python 3.7

Numpy methods do not handle IntEnum type correctly.

```python
from enum import IntEnum
import numpy as np
from numba import njit

class Matrix(IntEnum):
    r = 4
    c = 4
    
@njit
def test():
    return np.zeros((Matrix.r, Matrix.c), dtype='i8')

test()
```
Output:

```
TypingError                               Traceback (most recent call last)
<ipython-input-9-fbd55f77ab7c> in <module>
----> 1 test()

~/miniconda3/lib/python3.7/site-packages/numba/core/dispatcher.py in _compile_for_args(self, *args, **kws)
    412                 e.patch_message(msg)
    413 
--> 414             error_rewrite(e, 'typing')
    415         except errors.UnsupportedError as e:
    416             # Something unsupported is present in the user code, add help info

~/miniconda3/lib/python3.7/site-packages/numba/core/dispatcher.py in error_rewrite(e, issue_type)
    355                 raise e
    356             else:
--> 357                 raise e.with_traceback(None)
    358 
    359         argtypes = []

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function zeros>) found for signature:
 
 >>> zeros(UniTuple(IntEnum<int64>(Matrix) x 2), dtype=Literal[str](i8))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload of function 'zeros': File: numba/core/typing/npydecl.py: Line 507.
    With argument(s): '(UniTuple(IntEnum<int64>(Matrix) x 2), dtype=unicode_type)':
   No match.

During: resolving callee type: Function(<built-in function zeros>)
During: typing of call at <ipython-input-8-0e7fb952e8a4> (11)


File "<ipython-input-8-0e7fb952e8a4>", line 11:
def test():
    return np.zeros((Matrix.r, Matrix.c), dtype='i8')
```

Hints from issue discussion:
Thanks for the report. This is currently unsupported so labelled as a `feature_request`. Casting the `IntEnum` to `intp` should suffice as a workaround:
```python
from enum import IntEnum
import numpy as np
from numba import njit

class Matrix(IntEnum):
    r = 4
    c = 4

@njit
def test():
    return np.zeros((np.intp(Matrix.r), np.intp(Matrix.c)), dtype='int64')

z = test()
print(z, z.dtype)
```
Also, the `dtype` as a literal string has only been supported since 0.52 (latest release). I note there's a bug in it in that `i8` is not correctly handled, you need to spell out a type which is an attribute on the `numpy` top level module, e.g. `int64` from `np.int64`, https://github.com/numba/numba/issues/6632 tracks.

Commit message: Added test

Patch to review:
diff --git a/numba/core/typing/npydecl.py b/numba/core/typing/npydecl.py
index f3b89de52ac..5fe56ebd05e 100644
--- a/numba/core/typing/npydecl.py
+++ b/numba/core/typing/npydecl.py
@@ -455,7 +455,7 @@ def parse_shape(shape):
     if isinstance(shape, types.Integer):
         ndim = 1
     elif isinstance(shape, (types.Tuple, types.UniTuple)):
-        if all(isinstance(s, types.Integer) for s in shape):
+        if all(isinstance(s, (types.Integer, types.IntEnumMember)) for s in shape):
             ndim = len(shape)
     return ndim
 
diff --git a/numba/tests/test_enums.py b/numba/tests/test_enums.py
index 5f5a31e9d7b..535c3300007 100644
--- a/numba/tests/test_enums.py
+++ b/numba/tests/test_enums.py
@@ -162,6 +162,19 @@ def pyfun(x):
         for member in IntEnumWithNegatives:
             self.assertPreciseEqual(pyfun(member), cfunc(member))
 
+    def test_int_shape_cast(self):
+        def pyfun_empty(x):
+            return np.empty((x, x), dtype='int64').fill(-1)
+        def pyfun_zeros(x):
+            return np.zeros((x, x), dtype='int64')
+        def pyfun_ones(x):
+            return np.ones((x, x), dtype='int64')
+        for pyfun in [pyfun_empty, pyfun_zeros, pyfun_ones]:
+            cfunc = jit(nopython=True)(pyfun)
+            for member in IntEnumWithNegatives:
+                if member >= 0:
+                    self.assertPreciseEqual(pyfun(member), cfunc(member))
+
 
 if __name__ == '__main__':
     unittest.main()


Code review comments:
diff_hunk: @@ -455,7 +455,7 @@ def parse_shape(shape):
     if isinstance(shape, types.Integer):
         ndim = 1
     elif isinstance(shape, (types.Tuple, types.UniTuple)):
-        if all(isinstance(s, types.Integer) for s in shape):
+        if all(isinstance(s, (types.Integer, types.IntEnumMember)) for s in shape):
path: numba/core/typing/npydecl.py
text: @user1:
Although this file is not automatically checked by flake8 (it should be, but needs fixing and removing from the ignore list) this line is over-long (more than 80 chars) - perhaps this could be resolved e.g. by:

```suggestion
        int_tys = (types.Integer, types.IntEnumMember)
        if all(isinstance(s, int_tys) for s in shape):
```

@author:
Sounds good, I'll add that!