如何向函数传递1或2个np.array且不对单个数组做维度拆包？

你遇到的单Numpy数组被沿第一维度拆包的问题，本质是Python的*拆包规则会把所有可迭代对象（包含Numpy数组）按迭代维度拆分。如果你是在参数定义时使用*args，调用时直接传数组对象不加*，原本是不会自动拆包的，如果出现拆包基本都是调用时不小心给数组前加了*运算符导致的。

这里提供几个可替代*args的实现方案，或者兼容*args的修复方法：

方案1：固定参数加默认值（最推荐，仅需支持1-2个数组的场景） #

直接明确定义参数，第二个参数设默认值为None，完全不需要用可变参数，从根源避免拆包问题：

import numpy as np

def arr_transform(arr1, arr2=None):
    # 单数组处理逻辑
    if arr2 is None:
        # 示例：数组元素乘2
        return arr1 * 2
    # 双数组处理逻辑
    else:
        # 示例：两个数组对应元素相加
        return arr1 + arr2

调用方式：

• 单数组传入：arr_transform(np.array([1,2,3]))
• 双数组传入：arr_transform(np.array([1,2]), np.array([3,4]))

方案2：保留*args，增加参数校验逻辑 #

如果后续需要扩展支持更多数组参数，可以保留*args设计，在函数开头加参数校验和异常处理：

def arr_transform(*args):
    # 过滤筛选所有Numpy数组参数
    arr_list = [arg for arg in args if isinstance(arg, np.ndarray)]
    # 校验参数数量
    if len(arr_list) not in (1, 2):
        raise ValueError("仅支持传入1个或2个Numpy数组作为参数")
    
    arr1 = arr_list[0]
    arr2 = arr_list[1] if len(arr_list) == 2 else None
    # 后续你的转换处理逻辑
    ...

如果确实存在调用时误加*导致单个数组被拆成多个一维数组的场景，可以加一层自动合并逻辑：

def arr_transform(*args):
    # 自动识别被误拆的二维数组
    if all(isinstance(i, np.ndarray) and i.ndim == 1 for i in args):
        merged_arr = np.vstack(args)
        # 处理合并后的单数组
        return merged_arr * 2
    # 正常参数逻辑
    arr_list = [arg for arg in args if isinstance(arg, np.ndarray)]
    ...

方案3：关键字-only参数避免传参混淆 #

如果想要完全避免参数位置传错的问题，可以把第二个参数设置为仅支持关键字传递：

def arr_transform(arr1, /, *, arr2=None):
    # / 之前的参数仅支持按位置传递，* 之后的参数仅支持按关键字传递
    if arr2 is None:
        return arr1 * 2
    return arr1 + arr2

调用方式：

• 单数组传入：arr_transform(np.array([1,2,3]))
• 双数组传入：arr_transform(np.array([1,2]), arr2=np.array([3,4]))

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