Numba Wrapper Address Protocol (WAP) is broken

[guilhermeleobas]

Reporting a bug

• I have tried using the latest released version of Numba (most recent is
  visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
• I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

Reproducer is from Numba docs:

import numba, ctypes, ctypes.util, math

libm = ctypes.cdll.LoadLibrary(ctypes.util.find_library('m'))

class LibMCos(numba.types.WrapperAddressProtocol):
    def __wrapper_address__(self):
        return ctypes.cast(libm.cos, ctypes.c_voidp).value
    def signature(self):
        return numba.float64(numba.float64)

cos = LibMCos()

@numba.njit
def foo(f, x):
    return f(x)

print(foo(cos, 0.0))

$ python repro.py
Traceback (most recent call last):
  File "/home/guilhermeleobas/git/numba/repro.py", line 17, in <module>
    print(foo(cos, 0.0))
  File "/home/guilhermeleobas/git/numba/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/guilhermeleobas/git/numba/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
During: typing of argument at /home/guilhermeleobas/git/numba/repro.py (15)

File "repro.py", line 15:
def foo(f, x):
    return f(x)

[stuartarchibald]

Thanks for the report. This stopped working sometime between 0.53 and 0.54.

[stuartarchibald]

git bisect suggests: 98f733f is the first bad commit.

[stuartarchibald]

Given the contents of 98f733f, I wonder if WAP relied on some import ordering/side effect.

[stuartarchibald]

This "fixes" it:

diff --git a/numba/core/cpu.py b/numba/core/cpu.py
--- a/numba/core/cpu.py
+++ b/numba/core/cpu.py
@@ -58,6 +58,17 @@ class CPUContext(BaseContext):
         # Initialize NRT runtime
         rtsys.initialize(self)
 
+        from numba.core.types import FunctionType, WrapperAddressProtocol
+        from numba.extending import typeof_impl
+        @typeof_impl.register(WrapperAddressProtocol)
+        def typeof_function_type(val, c):
+            if isinstance(val, WrapperAddressProtocol):
+                sig = val.signature()
+            else:
+                raise NotImplementedError(
+                    f'function type from {type(val).__name__}')
+            return FunctionType(sig)
+
         # Add lower_extension attribute
         self.lower_extensions = {}
         from numba.parfors.parfor_lowering import _lower_parfor_parallel

the issue is essentially that the WrapperAddressProtocol is not registered early enough for typing to know what it is.

[guilhermeleobas]

Thanks, Stuart. I think the tests for WAP are not failing on CI because test.SubTest is swallowing any exception thrown by the compiler.