如何在Numba重载中使用*args替代显式函数签名

问题 #

需要为169个带有Numba重载的函数（含45种独特签名）添加统一的额外逻辑（示例为结果为0时触发警告），但使用*args的通用装饰器会触发Numba类型错误，手动为每个签名编写装饰器效率极低，尝试动态生成签名又陷入递归循环。

解决方案 #

通过动态生成与原函数签名完全匹配的包装函数，让Numba能正确识别重载签名。利用inspect.signature自动提取参数列表，结合exec生成适配的包装函数，同时通过局部变量避免递归问题。

完整实现代码 #

from numba import njit, extending, objmode
import warnings
import inspect

def universal_warn(func):
    # 提取原函数的参数列表字符串
    sig = inspect.signature(func)
    param_str = ", ".join(sig.parameters.keys())
    
    # 动态生成与原函数签名一致的包装函数，用_original_func避免递归
    inner_def = f"""
def inner({param_str}):
    result = _original_func({param_str})
    if not result:
        with objmode:
            warnings.warn("Result is zero")
    return result
"""
    # 执行代码并返回生成的inner函数
    namespace = {"_original_func": func, "objmode": objmode, "warnings": warnings}
    exec(inner_def, globals(), namespace)
    return namespace["inner"]

# 示例函数及Numba重载
def func(a, b, c):
    return a + b + c

@extending.overload(func)
def func_overload(a, b, c):
    @universal_warn
    @extending.register_jitable
    def impl(a, b, c):
        return a + b - c
    return impl

# 测试调用
@njit
def func_nb(a, b, c):
    return func(a, b, c)

args = 1, 2, 3
print(func_nb(*args), func(*args))

关键说明 #

1. 签名适配：通过inspect.signature自动获取原函数的参数结构，生成的inner函数与原函数签名完全一致，Numba可正确匹配重载规则。
2. 递归规避：使用_original_func作为原函数的局部指代，避免直接使用func导致的递归调用问题。
3. 通用性：无需修改装饰器代码，直接用@universal_warn装饰任意签名的重载实现函数即可。

多签名适配示例 #

针对不同参数数量的函数，装饰器可直接复用：

def another_func(x, y):
    return x * y

@extending.overload(another_func)
def another_func_overload(x, y):
    @universal_warn
    @extending.register_jitable
    def impl(x, y):
        return x - y
    return impl

@njit
def another_func_nb(x, y):
    return another_func(x, y)

print(another_func_nb(2, 2), another_func(2, 2))  # 输出0 4，同时触发警告

内容的提问来源于stack exchange，提问作者Nin17