如何为接收函数参数的@njit装饰函数编写正确类型注解

结论先行 #

你已知的字符串形式Numba类型标注是官方支持的标准写法，不存在适配Python标准typing模块的更优方案——Numba的JIT编译类型系统和Python原生类型系统完全独立，标准库typing.Callable无法表达“可在Numba nopython模式下执行”的约束。

各方案对比说明 #

• 不推荐使用numba.core.registry.CPUDispatcher作为参数类型
  该类型是@nb.njit装饰后函数对象的运行时类型，仅能匹配已经完成JIT包装的函数，无法覆盖Numba支持的其他合法可调用入参（比如nopython模式下可自动内联编译的lambda、Numba CFunc对象），标注覆盖范围过窄，会把合法入参误判为类型错误。
• 字符串形式Numba签名是兼容性最好的运行时标注
  这类字符串标注是Numba编译阶段专门识别的类型标记，不会干扰Python运行时和常规静态检查逻辑。你之前使用的"nb.njit(callable)[[int, int], int]"是合法写法，更简洁的等价写法是直接引用Numba类型系统中的可调用类型：

  import numba as nb
  
  @nb.njit
  def foo_nb(fn: "nb.types.Callable[[int, int], int]") -> int:
      return fn(0, 1)
  

• 兼顾静态检查的适配写法
  目前mypy、pyright等静态检查工具无法识别Numba自定义类型规则是官方已知问题，可以通过typing.TYPE_CHECKING分支做双套标注适配：

  import typing as _T
  import numba as nb
  
  if _T.TYPE_CHECKING:
      # 静态检查阶段使用标准Callable做基础校验
      NbIntFunc = _T.Callable[[int, int], int]
  else:
      # 运行时、Numba编译阶段使用JIT可调用类型
      NbIntFunc = "nb.types.Callable[[int, int], int]"
  
  @nb.njit
  def foo_nb(fn: NbIntFunc) -> int:
      return fn(0, 1)
  

  该写法既可以让Numba在编译时正确校验可调用参数的签名，也能让静态检查工具拦截明显错误的入参（比如示例中传入整数1的场景）。

补充说明 #

示例中直接传入lambda x, y: x + y的写法在Numba 0.58及以上版本可正常运行：Numba会在nopython模式下自动编译符合语法要求的lambda表达式；如果传入未做JIT适配的普通Python函数，即使签名和Callable[[int,int],int]完全匹配，依然会触发TypingError，这类错误只能在JIT编译阶段触发，无法通过静态类型标注提前拦截。

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