通过Pyx扩展Numpy Random时遇Python对象转bitgen_t的Cython错误

错误原因 #

1. 类型转换问题：PyCapsule_GetPointer返回void*类型，Cython无法自动将其转换为bitgen_t*，必须显式强制类型转换。
2. 临时引用安全问题：原代码中capsule = x.capsule创建的是Python临时对象，直接存储其C指针会导致引用计数异常，需将capsule声明为Cython的PyObject*类型以持有引用。

修改后的extension_rand.pyx代码 #

from libc.stdint cimport uint32_t
from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
from cpython.object cimport PyObject

import numpy as np
cimport numpy as np
cimport cython

from numpy.random cimport bitgen_t
from numpy.random import PCG64

np.import_array()

@cython.boundscheck(False)
@cython.wraparound(False)
def uniforms(Py_ssize_t n):
    """
    Create an array of `n` uniformly distributed doubles.
    A 'real' distribution would want to process the values into
    some non-uniform distribution
    """
    cdef Py_ssize_t i
    cdef bitgen_t *rng
    cdef const char *capsule_name = "BitGenerator"
    cdef double[::1] random_values
    cdef np.ndarray randoms
    # 声明capsule为PyObject*，持有引用避免临时对象被回收
    cdef PyObject *capsule

    x = PCG64()
    capsule = x.capsule
    # 验证capsule有效性
    if not PyCapsule_IsValid(capsule, capsule_name):
        raise ValueError("Invalid pointer to bit generator state")
    # 显式强制转换void*到bitgen_t*
    rng = <bitgen_t*> PyCapsule_GetPointer(capsule, capsule_name)
    random_values = np.empty(n, dtype='float64')
    with x.lock, nogil:
        for i in range(n):
            # 调用BitGenerator的next_double方法
            random_values[i] = rng.next_double(rng.state)
    randoms = np.asarray(random_values)

    return randoms

修改后的setup.py（适配Numpy 1.26+目录结构） #

# setup.py
from setuptools import setup
from Cython.Build import cythonize
from setuptools.extension import Extension
import numpy as np
import os

# 获取Numpy随机数库的正确路径
lib_path = os.path.join(np.get_include(), 'numpy', 'random', 'lib')

extending = Extension(
    name="extension_rand", 
    sources=['extension_rand.pyx'],
    include_dirs=[np.get_include()],
    library_dirs=[lib_path],
    libraries=['npyrandom'],
    extra_compile_args=["-O3"]  # 可选：开启编译优化
)

setup(
    ext_modules=cythonize(
        extending, language_level="3", annotate=True
    )
)

编译验证 #

执行原命令即可完成编译：

python setup.py build_ext --inplace

关键修改说明 #

• 新增from cpython.object cimport PyObject，将capsule声明为cdef PyObject *capsule，确保Python对象引用被正确持有，避免临时引用被垃圾回收导致指针失效。
• 使用<bitgen_t*>对PyCapsule_GetPointer返回值做显式强制类型转换，解决Cython类型不匹配错误。
• 调整lib_path为os.path.join(np.get_include(), 'numpy', 'random', 'lib')，适配Numpy 1.26+的目录结构，确保能找到libnpyrandom库。

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