如何在Numba njit函数中使用类对象而不编译整个类？

最优方案：用工厂函数返回不同后端的实例（兼容Numba的jitclass + 普通类） #

核心思路是将System从类改为工厂函数，根据后端类型返回不同的实例：numpy后端返回Numba兼容的jitclass实例，其他后端返回普通类实例。这样用户无需修改任何调用代码，完全保留a.D的简洁接口，同时自动支持Numba njit加速。

代码实现 #

import numba as nb
import numpy as np
from abc import ABC, abstractmethod

# 定义抽象基类，统一接口规范
class BaseSystem(ABC):
    @property
    @abstractmethod
    def D(self):
        pass

    # 可添加公共方法，避免重复代码
    def reset(self):
        # 示例公共逻辑：重置属性（根据后端不同可重写）
        if isinstance(self, NumpySystem):
            self.D = np.ones_like(self.D)
        else:
            self._D = [[1,1],[1,1]]

# Numpy后端：用jitclass实现，天然兼容Numba
@nb.jitclass([('D', nb.float64[:, :])])
class NumpySystem(BaseSystem):
    def __init__(self):
        self.D = np.ones((2, 2))

# 其他后端：普通类实现，无需兼容Numba
class OtherSystem(BaseSystem):
    def __init__(self):
        self._D = [[1,1],[1,1]]

    @property
    def D(self):
        return self._D

# 工厂函数：统一入口，根据后端返回对应实例
def System(backend='numpy'):
    if backend == 'numpy':
        return NumpySystem()
    else:
        return OtherSystem()

用户代码无需修改，直接使用 #

@nb.njit()
def user_provided_function(a):
    result = a.D * 2
    return result

# Numpy后端：正常用njit加速
b = System(backend='numpy')
out = user_provided_function(b)
print(out)
# 输出：[[2. 2.]
#        [2. 2.]]

# 切换其他后端：接口完全一致，无需修改代码
c = System(backend='other')
print(c.D)
# 输出：[[1, 1], [1, 1]]

方案优势 #

1. 零用户代码改动：保留原有的System()实例化和a.D调用方式，用户无需重构代码即可切换后端
2. 天然兼容Numba：numpy后端的实例是jitclass，直接被njit函数识别，无类型报错
3. 清晰的代码隔离：不同后端的实现完全分离，便于维护和扩展新后端
4. 公共逻辑复用：通过抽象基类提取公共方法，避免重复代码

备选方案：用装饰器自动转换System实例为Numba兼容对象 #

如果不想修改System类的结构（比如原类有大量已有逻辑），可以通过自定义装饰器，自动将传入njit函数的System实例转换为Numba能识别的结构（如命名元组），同时保留用户的调用接口。

代码实现 #

import numba as nb
import numpy as np
from functools import wraps
from collections import namedtuple

class System():
    def __init__(self, backend='numpy'):
        self.backend = backend
        if backend == 'numpy':
            self.D = np.ones((2,2))
        else:
            self.D = [[1,1],[1,1]]

    # 定义转换为Numba兼容对象的方法
    def _as_jitable(self):
        if self.backend != 'numpy':
            raise ValueError("仅numpy后端支持JIT编译")
        SystemJitData = namedtuple('SystemJitData', ['D'])
        return SystemJitData(D=self.D)

# 自定义装饰器：自动处理System实例转换
def njit_with_system(*args, **kwargs):
    def decorator(func):
        numba_func = nb.njit(*args, **kwargs)(func)

        @wraps(func)
        def wrapper(*func_args, **func_kwargs):
            # 转换位置参数中的System实例
            new_args = [arg._as_jitable() if isinstance(arg, System) else arg for arg in func_args]
            # 转换关键字参数中的System实例
            new_kwargs = {k: v._as_jitable() if isinstance(v, System) else v for k, v in func_kwargs.items()}
            return numba_func(*new_args, **new_kwargs)
        return wrapper
    return decorator

用户代码只需替换装饰器 #

@njit_with_system()
def user_provided_function(a):
    result = a.D * 2
    return result

b = System(backend='numpy')
out = user_provided_function(b)
print(out)

方案优势 #

• 无需修改原有System类结构，适合已有大量代码的场景
• 用户仅需替换njit装饰器，调用逻辑完全不变

方案局限 #

• 转换后的命名元组是不可变对象，若用户函数中需要修改System属性则无法支持
• 非numpy后端调用jit函数会直接报错（符合预期，但需提前告知用户）

内容的提问来源于stack exchange，提问作者P. Egli