使用入口点注册扩展
第三方包通常既有面向用户的 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
文件。