Repository: numba/numba

Issue:
inspect_types() raising KeyError
The following example, in numba v0.21.0, will raise a KeyError:

```
from __future__ import print_function
from numba import jit
import numpy as np


@jit
def distance(fromLons, fromLats, toLons, toLats):
    fromLons = np.radians(fromLons)
    fromLats = np.radians(fromLats)
    toLons = np.radians(toLons)
    toLats = np.radians(toLats)

    # Haversine formula
    return (6371.0 * 2.0 * np.arcsin(np.sqrt(np.sin((toLats - fromLats)/2.0) ** 2
                          + np.cos(fromLats) * np.cos(toLats)
                        * np.sin((toLons - fromLons)/2.0) ** 2)))



if __name__ == '__main__':
    a = np.linspace(0, 360, 20)
    b = np.linspace(10, 55, 25)

    distance(a, b[:20], 1.75, 1.5)
    distance(a[:, None], b[None, :], 1.75, 1.5)

    print("distance() types:\n", distance.inspect_types())
```

The traceback:

```
bash-4.1$ python numba_scce.py 
distance (array(float64, 1d, C), array(float64, 1d, C), float64, float64)
--------------------------------------------------------------------------------
Traceback (most recent call last):
  File "numba_scce.py", line 27, in <module>
    print("distance() types:\n", distance.inspect_types())
  File "/home/broot/scratch/miniconda/lib/python2.7/site-packages/numba/dispatcher.py", line 194, in inspect_types
    print(res.type_annotation, file=file)
  File "/home/broot/scratch/miniconda/lib/python2.7/site-packages/numba/annotations/type_annotations.py", line 266, in __str__
    return self.annotate()
  File "/home/broot/scratch/miniconda/lib/python2.7/site-packages/numba/annotations/type_annotations.py", line 127, in annotate
    groupedinst = self.prepare_annotations()
  File "/home/broot/scratch/miniconda/lib/python2.7/site-packages/numba/annotations/type_annotations.py", line 111, in prepare_annotations
    atype = self.typemap[inst.target.name]
KeyError: 'fromLons_1'
```

I should note that commenting out the first call to `distance()` will not result in an error.



Discussion:
I can reproduce a somewhat similar error since v0.22 or 0.23:

``` python
Traceback (most recent call last):
  ...
  File "...", line 477, in <module>
    @numba.jit("void(float64[:], float64[:], float64, float64, float64[:])", nopython=True, cache=True)
  File "C:\Anaconda3\lib\site-packages\numba\decorators.py", line 169, in wrapper disp.compile(sig)
  File "C:\Anaconda3\lib\site-packages\numba\dispatcher.py", line 508, in compile self._cache.save_overload(sig, cres)
  File "C:\Anaconda3\lib\site-packages\numba\dispatcher.py", line 884, in save_overload self._save_data(data_name, cres)
  File "C:\Anaconda3\lib\site-packages\numba\dispatcher.py", line 981, in _save_data data = cres._reduce()
  File "C:\Anaconda3\lib\site-packages\numba\compiler.py", line 73, in _reduce typeann = str(self.type_annotation)
  File "C:\Anaconda3\lib\site-packages\numba\annotations\type_annotations.py", line 264,
  in __str__ return self.annotate()
  File "C:\Anaconda3\lib\site-packages\numba\annotations\type_annotations.py", line 125, 
  in annotate groupedinst = self.prepare_annotations()
  File "C:\Anaconda3\lib\site-packages\numba\annotations\type_annotations.py", line 106,
  in prepare_annotations atype = self.typemap[inst.target.name]

KeyError: '_141_11'
```

When debugging the interesting lines of code in `type_annotations.py` I saw that the key is somehow present in the typemap, but the formatting of the `inst.targets` name is wrong:

``` python
>>> self.typemap['$141.41']
float64
>>> inst.target.name
'_141_11'
>>> self.typemap[inst.target.name]
Traceback (most recent call last):
  File "<debug input>", line 1, in <module>
KeyError: '_141_11'
```

Currently it seems like the error only occurs when the cache option is enabled.

I'm also getting this error when I enable the cache option for some functions. I'm digging around, trying to figure out what about the functions appears to lead to this.

Curious... I had a 3x10 array, J, and I was performing the following vectorized ops:

``` python
u = a*(J[0]*z - dx*J[2])
v = a*(J[1]*z - dy*J[2])
```

I stopped getting KeyErrors when I put the ops into a loop:

``` python
u = empty(10)
v = empty(10)
for i in range(10):
    u[i] = a*(J[0, i]*z - dx*J[2, i])
    v[i] = a*(J[1, i]*z - dy*J[2, i])
```

I have also hit this bug (similar context as to what @kitizz has), still present in 0.28.1 in Anaconda. `cache=False` indeed does skip this bug.

These are really two different bugs. The original bug (as posted by the OP) occurs even with caching disabled. Please post a new issue for the caching-specific bug.


Resolving patch:
diff --git a/numba/ir.py b/numba/ir.py
--- a/numba/ir.py
+++ b/numba/ir.py
@@ -148,6 +148,8 @@ def __init__(self, op, loc, **kws):
         self._kws = kws
 
     def __getattr__(self, name):
+        if name.startswith('_'):
+            return Inst.__getattr__(self, name)
         return self._kws[name]
 
     def __setattr__(self, name, value):
diff --git a/numba/npyufunc/array_exprs.py b/numba/npyufunc/array_exprs.py
--- a/numba/npyufunc/array_exprs.py
+++ b/numba/npyufunc/array_exprs.py
@@ -2,6 +2,7 @@
 
 import ast
 from collections import defaultdict, OrderedDict
+import contextlib
 import sys
 
 import numpy as np
@@ -302,46 +303,73 @@ def _arr_expr_to_ast(expr):
         "Don't know how to translate array expression '%r'" % (expr,))
 
 
+@contextlib.contextmanager
+def _legalize_parameter_names(var_list):
+    """
+    Legalize names in the variable list for use as a Python function's
+    parameter names.
+    """
+    var_map = OrderedDict()
+    for var in var_list:
+        old_name = var.name
+        new_name = old_name.replace("$", "_").replace(".", "_")
+        # Caller should ensure the names are unique
+        assert new_name not in var_map
+        var_map[new_name] = var, old_name
+        var.name = new_name
+    param_names = list(var_map)
+    try:
+        yield param_names
+    finally:
+        # Make sure the old names are restored, to avoid confusing
+        # other parts of Numba (see issue #1466)
+        for var, old_name in var_map.values():
+            var.name = old_name
+
+
 def _lower_array_expr(lowerer, expr):
     '''Lower an array expression built by RewriteArrayExprs.
     '''
     expr_name = "__numba_array_expr_%s" % (hex(hash(expr)).replace("-", "_"))
+    expr_filename = expr.loc.filename
     expr_var_list = expr.list_vars()
-    expr_var_map = {}
-    for expr_var in expr_var_list:
-        expr_var_name = expr_var.name
-        expr_var_new_name = expr_var_name.replace("$", "_").replace(".", "_")
-        # Avoid inserting existing var into the expr_var_map
-        if expr_var_new_name not in expr_var_map:
-            expr_var_map[expr_var_new_name] = expr_var_name, expr_var
-        expr_var.name = expr_var_new_name
-    expr_filename = expr_var_list[0].loc.filename
-    # Parameters are the names internal to the new closure.
-    expr_params = sorted(expr_var_map.keys())
-    # Arguments are the names external to the new closure (except in
-    # Python abstract syntax, apparently...)
-    expr_args = [expr_var_map[key][0] for key in expr_params]
-    if hasattr(ast, "arg"):
-        # Should be Python 3.x
-        ast_args = [ast.arg(param_name, None)
-                    for param_name in expr_params]
-    else:
-        # Should be Python 2.x
-        ast_args = [ast.Name(param_name, ast.Param())
-                    for param_name in expr_params]
-    # Parse a stub function to ensure the AST is populated with
-    # reasonable defaults for the Python version.
-    ast_module = ast.parse('def {0}(): return'.format(expr_name),
-                           expr_filename, 'exec')
-    assert hasattr(ast_module, 'body') and len(ast_module.body) == 1
-    ast_fn = ast_module.body[0]
-    ast_fn.args.args = ast_args
-    ast_fn.body[0].value, namespace = _arr_expr_to_ast(expr.expr)
-    ast.fix_missing_locations(ast_module)
+    # The expression may use a given variable several times, but we
+    # should only create one parameter for it.
+    expr_var_unique = sorted(set(expr_var_list), key=lambda var: var.name)
+
+    # Arguments are the names external to the new closure
+    expr_args = [var.name for var in expr_var_unique]
+
+    # 1. Create an AST tree from the array expression.
+
+    with _legalize_parameter_names(expr_var_unique) as expr_params:
+
+        if hasattr(ast, "arg"):
+            # Should be Python 3.x
+            ast_args = [ast.arg(param_name, None)
+                        for param_name in expr_params]
+        else:
+            # Should be Python 2.x
+            ast_args = [ast.Name(param_name, ast.Param())
+                        for param_name in expr_params]
+        # Parse a stub function to ensure the AST is populated with
+        # reasonable defaults for the Python version.
+        ast_module = ast.parse('def {0}(): return'.format(expr_name),
+                               expr_filename, 'exec')
+        assert hasattr(ast_module, 'body') and len(ast_module.body) == 1
+        ast_fn = ast_module.body[0]
+        ast_fn.args.args = ast_args
+        ast_fn.body[0].value, namespace = _arr_expr_to_ast(expr.expr)
+        ast.fix_missing_locations(ast_module)
+
+    # 2. Compile the AST module and extract the Python function.
+
     code_obj = compile(ast_module, expr_filename, 'exec')
     six.exec_(code_obj, namespace)
     impl = namespace[expr_name]
 
+    # 3. Now compile a ufunc using the Python function as kernel.
+
     context = lowerer.context
     builder = lowerer.builder
     outer_sig = expr.ty(*(lowerer.typeof(name) for name in expr_args))


Test patch:
diff --git a/numba/tests/test_array_exprs.py b/numba/tests/test_array_exprs.py
--- a/numba/tests/test_array_exprs.py
+++ b/numba/tests/test_array_exprs.py
@@ -6,7 +6,7 @@
 
 from numba import njit, vectorize
 from numba import unittest_support as unittest
-from numba import compiler, typing, typeof, ir
+from numba import compiler, typing, typeof, ir, utils
 from numba.compiler import Pipeline, _PipelineManager, Flags
 from numba.targets import cpu
 from .support import MemoryLeakMixin, TestCase
@@ -55,6 +55,12 @@ def explicit_output(a, b, out):
     np.cos(a, out)
     return np.add(out, b, out)
 
+def variable_name_reuse(a, b, c, d):
+    u = a + b
+    u = u - a * b
+    u = u * c + d
+    return u
+
 
 # From issue #1264
 def distance_matrix(vectors):
@@ -446,6 +452,22 @@ def foo(a, b):
         got = foo(a, b)
         self.assertPreciseEqual(got, expect)
 
+    def test_annotations(self):
+        """
+        Type annotation of array expressions with disambiguated
+        variable names (issue #1466).
+        """
+        cfunc = njit(variable_name_reuse)
+
+        a = np.linspace(0, 1, 10)
+        cfunc(a, a, a, a)
+
+        buf = utils.StringIO()
+        cfunc.inspect_types(buf)
+        res = buf.getvalue()
+        self.assertIn("#   u.1 = ", res)
+        self.assertIn("#   u.2 = ", res)
+
 
 class TestSemantics(MemoryLeakMixin, unittest.TestCase):