[Python] how to force numba to return a numpy type?

I find this behavior quite counter-intuitive although I suppose there is a reason for it - numba automatically converts my numpy integer types directly into a python int:

import numba as nb
import numpy as np

print(f"Numba version: {nb.__version__}") # 0.59.0
print(f"NumPy version: {np.__version__}") # 1.23.5

# Explicitly define the signature
sig = nb.uint32(nb.uint32, nb.uint32)

@nb.njit(sig, cache=False)
def test_fn(a, b):
return a * b

res = test_fn(2, 10)
print(f"Result value: {res}") # returns 20
print(f"Result type: {type(res)}") # returns <class 'int'>


This is an issue as I'm using the return as an input into another njit function so I get a casting warning (and I also do unnecessary casts in-between the njit functions)

Is there any way to force numba to give me np.uint32 as a result instead?

--- EDIT ---

This is the best I've managed to do myself, however I refuse to believe this is the best implementation out there:

# we manually define a return record and pass it as a parameter
res_type = np.dtype([('res', np.uint32)])
sig = nb.void(nb.uint32, nb.uint32, nb.from_dtype(res_type))

@nb.njit(sig, cache=False)
def test_fn(a:np.uint32, b:np.uint32, res: res_type):
res['res'] = a * b

# Call with Python ints (Numba should coerce based on signature)
res = np.recarray(1, dtype=res_type)[0]
res_py_in = test_fn(2, 10, res)
print(f"\nCalled with Python ints:")
print(f"Result value: {res['res']}") # 20
print(f"Result type: {type(res['res'])}") # <class 'numpy.uint32'>

Continue reading...