使用入口点注册扩展

第三方包通常既有面向用户的 API,也会定义 Numba 编译器的扩展。在这些情况下,当用户导入包时,新的类型和重载可以注册到 Numba。然而,在某些情况下,Numba 扩展通常不会被用户直接导入,但仍必须注册到 Numba 编译器。一个例子是 numba-scipy 包,它为 Numba 添加了对某些 SciPy 函数的支持。最终用户无需 import numba_scipy 即可启用对 SciPy 的编译器支持,该扩展只需安装在 Python 环境中即可。

Numba 使用 setuptools入口点 功能来发现扩展。这允许 Python 包注册一个初始化函数,该函数将在 numba 首次编译之前被调用。这种延迟确保了导入扩展的成本被推迟到必要时才发生。

添加对“Init”入口点的支持

包可以通过在 setup.py 中的 setup() 函数调用中添加 entry_points 参数来向 Numba 注册初始化函数。

setup(
    ...,
    entry_points={
        "numba_extensions": [
            "init = numba_scipy:_init_extension",
        ],
    },
    ...
)

Numba 当前只在 numba_extensions 组中查找 init 入口点。该入口点应为一个不接受任何参数的函数(任何名称,只要与 setup.py 中列出的匹配即可),其返回值将被忽略。此函数应注册类型、重载,或调用其他 Numba 扩展 API。扩展的初始化顺序是未定义的。

测试您的入口点

Numba 在第一个函数被编译时加载所有入口点。要测试您的入口点,仅 import numba 是不够的;您必须定义并运行一个小型函数,如下所示

import numba; numba.njit(lambda x: x + 1)(123)

无需导入您的模块:入口点是通过库的 *.egg-info 目录中的 entry_points.txt 文件识别的。

setup.py build 命令不会创建 egg,但 setup.py sdist(用于在本地目录中测试)和 setup.py install 会创建。所有在 Python 路径上的 egg 中注册的入口点都会被加载。调试时请务必检查是否存在过时的 entry_points.txt 文件。