如何用C函数复制Python二维Numpy数组并返回正确形状？

一、返回数组丢失第二维度的原因及修复 #

原因分析 #

你创建输出数组时未明确指定二维形状参数，仅按一维数组初始化，导致Numpy解析时丢失第二维度信息。

修复步骤 #

1. 获取输入数组的二维形状：通过Numpy C API的PyArray_DIMS获取行数和列数；
2. 创建二维输出数组：调用PyArray_SimpleNew时传入二维维度数组；
3. 按二维内存布局复制元素：确保元素访问符合Numpy的行/列优先存储规则。

二、matrix_out[i][j] = matrix_in[i][j];编译报错原因 #

Numpy数组在C层的内存是连续的一维块（即使是二维数组），并非T**类型的二维指针结构。直接用[i][j]二维下标访问会被编译器识别为非法的指针操作，因此报错。

正确的元素访问方式 #

有两种安全的访问方式：

1. 手动计算内存偏移：根据行优先/列优先规则计算元素在一维内存中的位置；
2. 使用Numpy官方宏：PyArray_GETPTR2可直接获取二维数组指定位置的元素指针。

完整修复代码示例 #

修改后的matrix_copy.c #

#include <Python.h>
#include <numpy/arrayobject.h>
#include <string.h>

static PyObject* copy_matrix(PyObject* self, PyObject* args) {
    PyArrayObject *matrix_in;
    // 解析输入参数，确保是Numpy数组
    if (!PyArg_ParseTuple(args, "O!", &PyArray_Type, &matrix_in)) {
        return NULL;
    }

    // 校验输入为二维数组
    if (PyArray_NDIM(matrix_in) != 2) {
        PyErr_SetString(PyExc_ValueError, "输入必须是二维Numpy数组");
        return NULL;
    }

    // 获取数组维度与类型信息
    npy_intp rows = PyArray_DIMS(matrix_in)[0];
    npy_intp cols = PyArray_DIMS(matrix_in)[1];
    int dtype = PyArray_TYPE(matrix_in);
    int item_size = PyArray_ITEMSIZE(matrix_in);

    // 创建二维输出数组（关键：指定2个维度）
    npy_intp dims[2] = {rows, cols};
    PyArrayObject *matrix_out = (PyArrayObject*)PyArray_SimpleNew(2, dims, dtype);
    if (!matrix_out) {
        return NULL;
    }

    // 获取数组数据的一维指针
    char *in_data = PyArray_DATA(matrix_in);
    char *out_data = PyArray_DATA(matrix_out);

    // 按行优先复制所有元素
    for (npy_intp i = 0; i < rows; i++) {
        for (npy_intp j = 0; j < cols; j++) {
            // 计算元素偏移量，复制对应字节
            size_t offset = (i * cols + j) * item_size;
            memcpy(out_data + offset, in_data + offset, item_size);
        }
    }

    return (PyObject*)matrix_out;
}

// 模块方法定义
static PyMethodDef MatrixCopyMethods[] = {
    {"copy_matrix", copy_matrix, METH_VARARGS, "复制二维Numpy数组并保留维度"},
    {NULL, NULL, 0, NULL}
};

// 模块定义
static struct PyModuleDef matrixcopymodule = {
    PyModuleDef_HEAD_INIT,
    "matrix_copy",
    NULL,
    -1,
    MatrixCopyMethods
};

// 模块初始化（必须调用import_array）
PyMODINIT_FUNC PyInit_matrix_copy(void) {
    import_array();
    return PyModule_Create(&matrixcopymodule);
}

测试用main.py #

import numpy as np
import matrix_copy

# 创建300x200的测试数组
test_arr = np.random.rand(300, 200).astype(np.float32)
copied_arr = matrix_copy.copy_matrix(test_arr)

print(f"输入数组形状: {test_arr.shape}")
print(f"输出数组形状: {copied_arr.shape}")
print(f"数组内容是否一致: {np.array_equal(test_arr, copied_arr)}")

编译命令（以Python3.8为例） #

gcc -shared -fPIC -I/usr/include/python3.8 -I/usr/lib/python3.8/site-packages/numpy/core/include matrix_copy.c -o matrix_copy.so

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