Repository: numba/numba

Issue:
Support getitem for type types.
The syntax for an array uses square brackets, for example: `int64[:]`. Numba currently does not support this syntax within an `njit`context. For example the following will not work:

```python
from numba import njit
from numba.types import int64
from numba.typed import List
import numpy as np


@njit
def foo():
    l = List.empty_list(int64[:])
    l.append(np.zeros(23))
    return l


foo()
```

And will instead error out with:

```
Traceback (most recent call last):
  File "/Users/vhaenel/git/numba/numba/core/typeinfer.py", line 155, in propagate
    constraint(typeinfer)
  File "/Users/vhaenel/git/numba/numba/core/typeinfer.py", line 398, in __call__
    self.fallback(typeinfer)
  File "/Users/vhaenel/git/numba/numba/core/typeinfer.py", line 618, in __call__
    self.resolve(typeinfer, typeinfer.typevars, fnty=fnty)
  File "/Users/vhaenel/git/numba/numba/core/typeinfer.py", line 515, in resolve
    sig = typeinfer.resolve_call(fnty, pos_args, kw_args)
  File "/Users/vhaenel/git/numba/numba/core/typeinfer.py", line 1459, in resolve_call
    return self.context.resolve_function_type(fnty, pos_args, kw_args)
  File "/Users/vhaenel/git/numba/numba/core/typing/context.py", line 214, in resolve_function_type
    raise last_exception
  File "/Users/vhaenel/git/numba/numba/core/typing/context.py", line 197, in resolve_function_type
    res = self._resolve_user_function_type(func, args, kws)
  File "/Users/vhaenel/git/numba/numba/core/typing/context.py", line 249, in _resolve_user_function_type
    return func.get_call_type(self, args, kws)
  File "/Users/vhaenel/git/numba/numba/core/types/functions.py", line 148, in get_call_type
    failures.raise_error()
  File "/Users/vhaenel/git/numba/numba/core/types/functions.py", line 77, in raise_error
    raise errors.TypingError(self.format())
numba.core.errors.TypingError: Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (class(int64), slice<a:b>)
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:
    All templates rejected with literals.
In definition 7:
    All templates rejected without literals.
In definition 8:
    All templates rejected with literals.
In definition 9:
    All templates rejected without literals.
In definition 10:
    All templates rejected with literals.
In definition 11:
    All templates rejected without literals.
In definition 12:
    All templates rejected with literals.
In definition 13:
    All templates rejected without literals.
In definition 14:
    All templates rejected with literals.
In definition 15:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: typing of intrinsic-call at foo070.py (9)
[2] During: typing of static-get-item at foo070.py (9)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "foo070.py", line 14, in <module>
    foo()
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 342, in error_rewrite
    raise e
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 353, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/Users/vhaenel/git/numba/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 794, in compile
    cres = self._compiler.compile(args, return_type)
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 81, in compile
    raise retval
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 91, in _compile_cached
    retval = self._compile_core(args, return_type)
  File "/Users/vhaenel/git/numba/numba/core/dispatcher.py", line 109, in _compile_core
    pipeline_class=self.pipeline_class)
  File "/Users/vhaenel/git/numba/numba/core/compiler.py", line 568, in compile_extra
    return pipeline.compile_extra(func)
  File "/Users/vhaenel/git/numba/numba/core/compiler.py", line 339, in compile_extra
    return self._compile_bytecode()
  File "/Users/vhaenel/git/numba/numba/core/compiler.py", line 401, in _compile_bytecode
    return self._compile_core()
  File "/Users/vhaenel/git/numba/numba/core/compiler.py", line 381, in _compile_core
    raise e
  File "/Users/vhaenel/git/numba/numba/core/compiler.py", line 372, in _compile_core
    pm.run(self.state)
  File "/Users/vhaenel/git/numba/numba/core/compiler_machinery.py", line 341, in run
    raise patched_exception
  File "/Users/vhaenel/git/numba/numba/core/compiler_machinery.py", line 332, in run
    self._runPass(idx, pass_inst, state)
  File "/Users/vhaenel/git/numba/numba/core/compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/vhaenel/git/numba/numba/core/compiler_machinery.py", line 291, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/Users/vhaenel/git/numba/numba/core/compiler_machinery.py", line 264, in check
    mangled = func(compiler_state)
  File "/Users/vhaenel/git/numba/numba/core/typed_passes.py", line 98, in run_pass
    raise_errors=self._raise_errors)
  File "/Users/vhaenel/git/numba/numba/core/typed_passes.py", line 70, in type_inference_stage
    infer.propagate(raise_errors=raise_errors)
  File "/Users/vhaenel/git/numba/numba/core/typeinfer.py", line 994, in propagate
    raise errors[0]
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (class(int64), slice<a:b>)
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:
    All templates rejected with literals.
In definition 7:
    All templates rejected without literals.
In definition 8:
    All templates rejected with literals.
In definition 9:
    All templates rejected without literals.
In definition 10:
    All templates rejected with literals.
In definition 11:
    All templates rejected without literals.
In definition 12:
    All templates rejected with literals.
In definition 13:
    All templates rejected without literals.
In definition 14:
    All templates rejected with literals.
In definition 15:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: typing of intrinsic-call at foo070.py (9)
[2] During: typing of static-get-item at foo070.py (9)

File "foo070.py", line 9:
def foo():
    l = List.empty_list(int64[:])
    ^
```
 

Implementing `getitem` for type types should resolve this.


Discussion:
Can I work on this?

Resolving patch:
diff --git a/numba/core/typing/builtins.py b/numba/core/typing/builtins.py
--- a/numba/core/typing/builtins.py
+++ b/numba/core/typing/builtins.py
@@ -654,6 +654,24 @@ def generic(self, args, kws):
             sig = signature(ret, *args)
             return sig
 
+@infer
+class StaticGetItemClass(AbstractTemplate):
+    """This handles the "static_getitem" when a Numba type is subscripted e.g:
+    var = typed.List.empty_list(float64[::1, :])
+    It only allows this on simple numerical types. Compound types, like
+    records, are not supported.
+    """
+    key = "static_getitem"
+
+    def generic(self, args, kws):
+        clazz, idx = args
+        if not isinstance(clazz, types.NumberClass):
+            return
+        ret = clazz.dtype[idx]
+        sig = signature(ret, *args)
+        return sig
+
+
 # Generic implementation for "not in"
 
 @infer
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
--- a/numba/np/arrayobj.py
+++ b/numba/np/arrayobj.py
@@ -2922,6 +2922,31 @@ def dtype_type(context, builder, dtypety, dtypeval):
 lower_getattr(types.DType, 'kind')(dtype_type)
 
 
+# ------------------------------------------------------------------------------
+# static_getitem on Numba numerical types to create "array" types
+
+
+@lower_builtin('static_getitem', types.NumberClass, types.Any)
+def static_getitem_number_clazz(context, builder, sig, args):
+    """This handles the "static_getitem" when a Numba type is subscripted e.g:
+    var = typed.List.empty_list(float64[::1, :])
+    It only allows this on simple numerical types. Compound types, like
+    records, are not supported.
+    """
+    retty = sig.return_type
+    if isinstance(retty, types.Array):
+        # This isn't used or practically accessible, but has to exist, so just
+        # put in a NULL of the right type.
+        res = context.get_value_type(retty)(None)
+        return impl_ret_untracked(context, builder, retty, res)
+    else:
+        # This should be unreachable unless the implementation on the Type
+        # metaclass is changed.
+        msg = ("Unreachable; the definition of __getitem__ on the "
+               "numba.types.abstract.Type metaclass should prevent access.")
+        raise errors.LoweringError(msg)
+
+
 # ------------------------------------------------------------------------------
 # Structured / record lookup
 


Test patch:
diff --git a/numba/tests/test_getitem_on_types.py b/numba/tests/test_getitem_on_types.py
new file mode 100644
--- /dev/null
+++ b/numba/tests/test_getitem_on_types.py
@@ -0,0 +1,107 @@
+import unittest
+from itertools import product
+from numba import types, njit, typed, errors
+from numba.tests.support import TestCase
+
+
+class TestGetitemOnTypes(TestCase):
+    # Tests getitem on the type types.
+
+    def test_static_getitem_on_type(self):
+
+        def gen(numba_type, index):
+            def foo():
+                ty = numba_type[index] # a static_getitem
+                return typed.List.empty_list(ty)
+            return foo
+
+        # test a few types
+        tys = (types.bool_, types.float64, types.uint8, types.complex128,)
+
+        # and a few indexes of increasing complexity
+        contig = slice(None, None, 1) # unit stride
+        noncontig = slice(None, None, None)
+        indexes = (contig, # 1d contig -> C order
+                   noncontig, # 1d non-contig -> A order
+                   (noncontig, contig), # 2d C order
+                   (contig, noncontig), # 2d F order
+                   (noncontig, noncontig), # 2d A order
+                   (noncontig, noncontig, contig), # 3d C order
+                   (contig, noncontig, noncontig), # 3d F order
+                   (noncontig, noncontig, noncontig), # 3d A order
+                   )
+
+        for ty, idx in product(tys, indexes):
+            compilable = njit(gen(ty, idx))
+            # check the type of the typed list returned matches the type
+            # as constructed in the interpreter
+            expected = ty[idx]
+            # check execution
+            self.assertEqual(compilable()._dtype, expected)
+            got = compilable.nopython_signatures[0].return_type.dtype
+            # check sig
+            self.assertEqual(got, expected)
+
+    def test_shorthand_syntax(self):
+        # tests a couple of shorthand syntax examples
+        # (test_static_getitem_on_type is a more extensive test of the
+        # functionality but it uses slices directly).
+
+        @njit
+        def foo1():
+            ty = types.float32[::1, :] # 2d F order
+            return typed.List.empty_list(ty)
+
+        self.assertEqual(foo1()._dtype, types.float32[::1, :])
+
+        @njit
+        def foo2():
+            ty = types.complex64[:, :, :] # 3d A order
+            return typed.List.empty_list(ty)
+
+        self.assertEqual(foo2()._dtype, types.complex64[:, :, :])
+
+    def test_static_getitem_on_invalid_type(self):
+        # check that an unsupported type cannot be instantiated in njit code
+
+        # check this executes in the interpreter:
+        types.void[:]
+
+        # check the same fails in compilation as it's not supported
+        # it'll fall back to a generic getitem
+        with self.assertRaises(errors.TypingError) as raises:
+            @njit
+            def foo():
+                types.void[:]
+
+            foo()
+
+        msg = ("No implementation",
+               "getitem(typeref[none], slice<a:b>)")
+
+        excstr = str(raises.exception)
+        for m in msg:
+            self.assertIn(m, excstr)
+
+    def test_standard_getitem_on_type(self):
+        # not supported at present, should be doable if the slice is a literal
+        # though.
+
+        # check using a non-static arg to the getitem raises
+        with self.assertRaises(errors.TypingError) as raises:
+            @njit
+            def foo(not_static):
+                types.float64[not_static]
+
+            foo(slice(None, None, 1))
+
+        msg = ("No implementation",
+               "getitem(class(float64), slice<a:b>)")
+
+        excstr = str(raises.exception)
+        for m in msg:
+            self.assertIn(m, excstr)
+
+
+if __name__ == '__main__':
+    unittest.main()