Repository: numba/numba

Issue:
literally fails when getitem by string literal index is involved
- [x] I am using the latest released version of Numba (most recent is visible in
 the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [x] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

Hi, 

We've encountered a problem with literally and getitem by string literal index trying to migrate to 0.49.0rc2. It seems that RewriteStringLiteralGetitems recently added in #5182  transforms getitem expressions in IR to static_getitem, which then fails during lowering if no overload for static_getitem was provided for the user types, and because falling back to generic getitem fails to find signature for the already rewritten expression, here:

from lowering.py:
```
        elif expr.op == "static_getitem":
            signature = typing.signature(
                resty,
                self.typeof(expr.value.name),
                _lit_or_omitted(expr.index),
            )
           ...
            except NotImplementedError:
                if expr.index_var is None:
                    raise
                # Fall back on the generic getitem() implementation
                # for this type.
                signature = self.fndesc.calltypes[expr]            # <---- this searches for 'static_getitem' expr
                return self.lower_getitem(resty, expr, expr.value,
                                          expr.index_var, signature)

```
which results in errors similar to:
```
Traceback (most recent call last):
  File "c:\users\akozlov\appdata\local\continuum\anaconda3\numba_master\numba\numba\core\errors.py", line 801, in new_error_context
    yield
  File "c:\users\akozlov\appdata\local\continuum\anaconda3\numba_master\numba\numba\core\lowering.py", line 267, in lower_block
    self.lower_inst(inst)
  File "c:\users\akozlov\appdata\local\continuum\anaconda3\numba_master\numba\numba\core\lowering.py", line 364, in lower_inst
    val = self.lower_assign(ty, inst)
  File "c:\users\akozlov\appdata\local\continuum\anaconda3\numba_master\numba\numba\core\lowering.py", line 538, in lower_assign
    return self.lower_expr(ty, value)
  File "c:\users\akozlov\appdata\local\continuum\anaconda3\numba_master\numba\numba\core\lowering.py", line 1183, in lower_expr
    signature = self.fndesc.calltypes[expr]
KeyError: static_getitem(value=$8call_method.3, index=C, index_var=col_name)
```

A following script can be used to reproduce the issue:
```
import operator
import numpy as np

import numba
from numba import types
from numba.core import cgutils
from numba.extending import (
    make_attribute_wrapper,
    overload,
    overload_method,
    models,
    register_model,
    type_callable,
    typeof_impl,
    lower_builtin,
)
from numba import literally

def make_dummy_type():
    class Dummy(object):
        def __init__(self, value):
            self.value = value

    class DummyType(types.Type):
        def __init__(self):
            super(DummyType, self).__init__(name="Dummy")

    dummy_type = DummyType()

    @register_model(DummyType)
    class DummyModel(models.StructModel):
        def __init__(self, dmm, fe_type):
            members = [("value", types.intp)]
            models.StructModel.__init__(self, dmm, fe_type, members)

    make_attribute_wrapper(DummyType, "value", "value")

    @type_callable(Dummy)
    def type_dummy(context):
        def typer(value):
            return dummy_type
        return typer

    @lower_builtin(Dummy, types.intp)
    def impl_dummy(context, builder, sig, args):
        typ = sig.return_type
        [value] = args
        dummy = cgutils.create_struct_proxy(typ)(context, builder)
        dummy.value = value
        return dummy._getvalue()

    @typeof_impl.register(Dummy)
    def typeof_dummy(val, c):
        return DummyType()

    # Store attributes
    return Dummy, DummyType

Dummy, DummyType = make_dummy_type()


@overload(operator.getitem)
def dummy_getitem_ovld(self, idx):
    if not isinstance(self, DummyType):
        return None

    print("DEBUG: typing Dummy::foo, idx:", idx)

    # suppose we can only support idx as literal argument
    if isinstance(idx, types.StringLiteral):
        def dummy_getitem_impl(self, idx):
            return 42
        return dummy_getitem_impl

    if isinstance(idx, types.UnicodeType):
        def dummy_getitem_impl(self, idx):
            return literally(idx)
        return dummy_getitem_impl

    return None

@numba.njit
def test_impl(x, y):
    return Dummy(x)[y]

if __name__ == "__main__":

    var = 'abc'
    print(test_impl(1, var))
```

Removing the line that updates signature or placing it in try catch+pass helps in this particular case, but I'm not really sure it's acceptable solution. So can you please take a look at this and find a proper solution for this problem, as it's essential for us to use 0.49.

Kind Regards,
Alexey.


Discussion:
Thanks for the report, I can reproduce.

Resolving patch:
diff --git a/numba/core/rewrites/static_getitem.py b/numba/core/rewrites/static_getitem.py
--- a/numba/core/rewrites/static_getitem.py
+++ b/numba/core/rewrites/static_getitem.py
@@ -61,6 +61,7 @@ def match(self, func_ir, block, typemap, calltypes):
         """
         self.getitems = getitems = {}
         self.block = block
+        self.calltypes = calltypes
         for expr in block.find_exprs(op='getitem'):
             if expr.op == 'getitem':
                 index_ty = typemap[expr.index.name]
@@ -85,6 +86,7 @@ def apply(self):
                                                       index=lit_val,
                                                       index_var=expr.index,
                                                       loc=expr.loc)
+                    self.calltypes[new_expr] = self.calltypes[expr]
                     inst = ir.Assign(value=new_expr, target=inst.target,
                                      loc=inst.loc)
             new_block.append(inst)


Test patch:
diff --git a/numba/tests/test_extending_types.py b/numba/tests/test_extending_types.py
--- a/numba/tests/test_extending_types.py
+++ b/numba/tests/test_extending_types.py
@@ -1,8 +1,9 @@
 """
 Test extending types via the numba.extending.* API.
 """
+import operator
 
-from numba import njit
+from numba import njit, literally
 from numba.core import types, cgutils
 from numba.core.errors import TypingError
 from numba.core.extending import lower_builtin
@@ -33,7 +34,7 @@ class DummyModel(models.StructModel):
             def __init__(self, dmm, fe_type):
                 members = [
                     ('value', types.intp),
-                    ]
+                ]
                 models.StructModel.__init__(self, dmm, fe_type, members)
 
         make_attribute_wrapper(DummyType, 'value', 'value')
@@ -128,3 +129,34 @@ def foo(x):
             foo(123)
         self.assertIn("cannot convert native Dummy to Python object",
                       str(raises.exception))
+
+    def test_issue5565_literal_getitem(self):
+        # the following test is adapted from
+        # https://github.com/numba/numba/issues/5565
+        Dummy, DummyType = self.Dummy, self.DummyType
+
+        MAGIC_NUMBER = 12321
+
+        @overload(operator.getitem)
+        def dummy_getitem_ovld(self, idx):
+            if not isinstance(self, DummyType):
+                return None
+            # suppose we can only support idx as literal argument
+            if isinstance(idx, types.StringLiteral):
+                def dummy_getitem_impl(self, idx):
+                    return MAGIC_NUMBER
+                return dummy_getitem_impl
+
+            if isinstance(idx, types.UnicodeType):
+                def dummy_getitem_impl(self, idx):
+                    return literally(idx)
+                return dummy_getitem_impl
+
+            return None
+
+        @njit
+        def test_impl(x, y):
+            return Dummy(x)[y]
+
+        var = 'abc'
+        self.assertEqual(test_impl(1, var), MAGIC_NUMBER)