sklam · Mar 4, 2021 · Feb 24, 2021 · Feb 25, 2021 · Mar 1, 2021 · stuartarchibald
diff --git a/numba/np/ufunc/array_exprs.py b/numba/np/ufunc/array_exprs.py
@@ -77,7 +77,10 @@ def _match_array_expr(self, instr, expr, target_name):
         if ((expr_op in ('unary', 'binop')) and (
                 expr.fn in npydecl.supported_array_operators)):
             # It is an array operator that maps to a ufunc.
-            array_assigns[target_name] = instr
+            # check that all args have internal types
+            if all(self.typemap[var.name].is_internal
+                   for var in expr.list_vars()):
+                array_assigns[target_name] = instr
 
         elif ((expr_op == 'call') and (expr.func.name in self.typemap)):
             # It could be a match for a known ufunc call.

diff --git a/numba/tests/test_array_exprs.py b/numba/tests/test_array_exprs.py
@@ -5,9 +5,27 @@
 
 from numba import njit, vectorize
 from numba import typeof
-from numba.core import utils, types, typing, ir, compiler, cpu
+from numba.core import utils, types, typing, ir, compiler, cpu, cgutils
 from numba.core.compiler import Compiler, Flags
-from numba.tests.support import MemoryLeakMixin, TestCase
+from numba.tests.support import MemoryLeakMixin, TestCase, temp_directory
+from numba.extending import (
+    overload,
+    models,
+    lower_builtin,
+    register_model,
+    make_attribute_wrapper,
+    type_callable,
+    typeof_impl
+)
+import operator
+import contextlib
+import importlib
+import shutil
+import uuid
+import os
+import sys
+import textwrap
+
 import unittest
 
 
@@ -627,5 +645,127 @@ def do_call(x, y):
         self.assertTrue(isinstance(oty.type.dtype, types.Float))
 
 
+class TestExternalTypes(MemoryLeakMixin, unittest.TestCase):
+    """ Tests RewriteArrayExprs with external (user defined) types,
+    see #5157"""
+
+    source_lines = textwrap.dedent("""
+        from numba.core import types
+
+        class FooType(types.Type):
+            def __init__(self):
+                super(FooType, self).__init__(name='Foo')
+        """)
+
+    # copied from test_types
+    @contextlib.contextmanager
+    def create_temp_module(self, source_lines=None, **jit_options):
+        # Use try/finally so cleanup happens even when an exception is raised
+        try:
+            if source_lines is None:
+                source_lines = self.source_lines
+            tempdir = temp_directory('test_extension_type')
+            # Generate random module name
+            temp_module_name = 'test_extension_type_{}'.format(
+                str(uuid.uuid4()).replace('-', '_'))
+            temp_module_path = os.path.join(tempdir, temp_module_name + '.py')
+
+            with open(temp_module_path, 'w') as f:
+                lines = source_lines.format(jit_options=jit_options)
+                f.write(lines)
+            # Add test_module to sys.path so it can be imported
+            sys.path.insert(0, tempdir)
+            test_module = importlib.import_module(temp_module_name)
+            yield test_module
+        finally:
+            sys.modules.pop(temp_module_name, None)
+            sys.path.remove(tempdir)
+            shutil.rmtree(tempdir)
+
+    def make_foo_type(self, FooType):
+        class Foo(object):
+            def __init__(self, value):
+                self.value = value
+
+        @register_model(FooType)
+        class FooModel(models.StructModel):
+            def __init__(self, dmm, fe_type):
+                members = [("value", types.intp)]
+                models.StructModel.__init__(self, dmm, fe_type, members)
+
+        make_attribute_wrapper(FooType, "value", "value")
+
+        @type_callable(Foo)
+        def type_foo(context):
+            def typer(value):
+                return FooType()
+
+            return typer
+
+        @lower_builtin(Foo, types.intp)
+        def impl_foo(context, builder, sig, args):
+            typ = sig.return_type
+            [value] = args
+            foo = cgutils.create_struct_proxy(typ)(context, builder)
+            foo.value = value
+            return foo._getvalue()
+
+        @typeof_impl.register(Foo)
+        def typeof_foo(val, c):
+            return FooType()
+
+        return Foo, FooType
+
+    def test_external_type(self):
+        with self.create_temp_module(self.source_lines) as test_module:
+            Foo, FooType = self.make_foo_type(test_module.FooType)
+
+            # sum of foo class instance and array return an array
+            # binary operation with foo class instance as one of args
+            @overload(operator.add)
+            def overload_foo_add(lhs, rhs):
+                if isinstance(lhs, FooType) and isinstance(rhs, types.Array):
+                    def imp(lhs, rhs):
+                        return np.array([lhs.value, rhs[0]])
+
+                    return imp
+
+            # sum of 2 foo class instances return an array
+            # binary operation with 2 foo class instances as args
+            @overload(operator.add)
+            def overload_foo_add(lhs, rhs):
+                if isinstance(lhs, FooType) and isinstance(rhs, FooType):
+                    def imp(lhs, rhs):
+                        return np.array([lhs.value, rhs.value])
+
+                    return imp
+
+            # neg of foo class instance return an array
+            # unary operation with foo class instance arg
+            @overload(operator.neg)
+            def overload_foo_neg(x):
+                if isinstance(x, FooType):
+                    def imp(x):
+                        return np.array([-x.value])
+
+                    return imp
+
+            @njit
+            def arr_expr_sum1(x, y):
+                return Foo(x) + np.array([y])
+
+            @njit
+            def arr_expr_sum2(x, y):
+                return Foo(x) + Foo(y)
+
+            @njit
+            def arr_expr_neg(x):
+                return -Foo(x)
+
+            np.testing.assert_array_equal(arr_expr_sum1(0, 1), np.array([0, 1]))
+            np.testing.assert_array_equal(arr_expr_sum2(2, 3), np.array([2, 3]))
+            np.testing.assert_array_equal(arr_expr_neg(4), np.array([-4]))
+
+
 if __name__ == "__main__":
     unittest.main()