如何在Numba jitted类中调用Cython扩展类型的方法？

这个问题的核心在于Numba的nopython JIT模式无法直接识别Cython扩展类型的方法——因为这些方法不属于Numba原生支持的类型系统范畴。下面我给你两种可行的修改方案，分别适配不同的场景需求：

方案一：用Numba的objmode做快速兼容 #

如果你的性能要求不是极端苛刻，最简单的方式是把调用Cython方法的代码放到Numba的objmode上下文管理器中。objmode会让Numba暂时回退到解释执行模式，处理它无法JIT编译的代码片段。

修改后的Python代码 #

import numba
from numba import jitclass, objmode
from shrubbery import Shrubbery

# JIT类的属性声明，shrub类型为object（兼容Cython扩展实例）
spec = [
    ('shrub', object)
]

@jitclass(spec)
class NumbaClass:
    def __init__(self, size):
        self.shrub = Shrubbery(size)
    
    def get_shrub_size(self):
        # 用objmode包裹Cython方法调用，指定返回值类型为intp（对应Python的int）
        with objmode(result='intp'):
            result = self.shrub.get_size()
        return result

# 测试运行
nc = NumbaClass(5)
print(nc.get_shrub_size())  # 输出：5

说明 #

• 这种方式几乎不需要修改你的Cython代码，只需要在Python端做少量调整。
• 缺点是objmode内部的代码会以解释执行的方式运行，如果这部分是性能热点，会有一定的开销。

方案二：导出Cython方法为C函数，让Numba直接调用 #

如果你的场景对性能要求很高，推荐把Cython的方法导出为可被Numba识别的C函数，通过C层面的调用实现零开销的JIT兼容。

步骤1：修改Cython代码（shrubbery.pyx） #

# 用public修饰类，确保它的类型信息能被外部C代码识别
cdef public class Shrubbery:
    cdef int size
    
    def __init__(self, size):
        self.size = size
    
    # 用cdef public导出内部方法，方便后续封装
    cdef public int _get_size(self):
        return self.size
    
    # 保留原有的Python方法，兼容普通Python调用
    def get_size(self):
        return self._get_size()

# 导出一个可以被C直接调用的函数，接收Shrubbery实例作为参数
cdef public int shrub_get_size(Shrubbery s):
    return s._get_size()

步骤2：修改Python代码，绑定C函数并适配Numba #

import numba
from numba import jitclass
from shrubbery import Shrubbery
import ctypes
import os

# 加载编译后的Cython扩展库（Windows下后缀为.pyd）
lib_path = os.path.abspath("shrubbery.so")
lib = ctypes.CDLL(lib_path)

# 定义C函数的参数和返回值类型
lib.shrub_get_size.argtypes = [ctypes.py_object]
lib.shrub_get_size.restype = ctypes.c_int

# 用Numba的cfunc把ctypes函数包装成可JIT调用的函数
@numba.cfunc(numba.int32(numba.pyobject))
def numba_shrub_get_size(s):
    return lib.shrub_get_size(s)

# JIT类的属性声明不变
spec = [
    ('shrub', object)
]

@jitclass(spec)
class NumbaClass:
    def __init__(self, size):
        self.shrub = Shrubbery(size)
    
    def get_shrub_size(self):
        # 直接调用包装后的Numba兼容函数
        return numba_shrub_get_size(self.shrub)

# 测试运行
nc = NumbaClass(5)
print(nc.get_shrub_size())  # 输出：5

说明 #

• 这种方式的性能几乎和纯Cython调用持平，因为Numba可以直接JIT编译到C函数的调用逻辑。
• 需要注意编译Cython时，public修饰的类和函数会被导出到符号表，确保ctypes能正确找到它们。

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