如何为返回指定尺寸Numpy数组的Python函数编写文档？

一、正确的类型标注方式 #

Python原生类型标注系统不直接支持指定NumPy数组的形状，但可以结合numpy.typing.NDArray和typing.Annotated补充形状信息，让VSCode这类编辑器识别更精准的类型提示：

from typing import Annotated
import numpy as np
from numpy.typing import NDArray

def poseToVector6(pose: Pose) -> Annotated[NDArray[np.float64], (6, 1)]:
    """将Pose转换为6×1的列向量"""
    # 函数实现
    return np.zeros((6, 1), dtype=np.float64)

如果不想引入额外的类型标注语法，也可以直接用NDArray并在文档字符串里明确形状：

import numpy as np
from numpy.typing import NDArray

def poseToVector6(pose: Pose) -> NDArray[np.float64]:
    """将Pose转换为6×1的列向量

    Returns:
        NDArray[np.float64]: 形状为(6, 1)的列向量，存储Pose的转换结果
    """
    # 函数实现
    return np.zeros((6, 1), dtype=np.float64)

二、用Docstring明确数组形状 #

详细的文档字符串是补充返回值信息的核心，推荐使用NumPy风格或Google风格的格式：

NumPy风格示例： #

import numpy as np

def poseToVector6(pose: Pose) -> np.ndarray:
    """将Pose对象转换为6×1的列向量

    Parameters
    ----------
    pose : Pose
        输入的Pose对象

    Returns
    -------
    np.ndarray
        形状为(6, 1)的float64类型列向量，存储Pose的平移和旋转分量
    """
    # 函数实现

Google风格示例： #

import numpy as np

def poseToVector6(pose: Pose) -> np.ndarray:
    """将Pose对象转换为6×1的列向量

    Args:
        pose: 输入的Pose对象

    Returns:
        np.ndarray: 形状为(6, 1)的列向量，类型为float64，包含平移(x,y,z)和旋转(r,p,y)分量
    """
    # 函数实现

三、常见误区说明 #

你之前尝试的-> np.ndarray((6,1))是错误写法：np.ndarray((6,1))是创建一个具体的数组实例，而非指定类型，因此类型检查器会将返回值识别为Any，这就是VSCode显示返回类型为Any的原因。

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