Python中使用itemgetter保留嵌套列表结构的方法

这个问题我之前也碰到过，operator.itemgetter的设计就是这样：当传入多个索引时返回元组，单个索引时直接返回对应元素，确实会导致后续处理的格式不统一。这里有几个靠谱的解决办法：

方法1：强制将结果转为元组 #

不管inds长度是多少，都把itemgetter的结果包装成元组。如果是单个元素，就手动构造成单元素元组：

from operator import itemgetter

a = [[1, 2, 3], [2, 3, 4], [1, 2], [3, 5]]
inds = [0]  # 单个索引的情况

getter = itemgetter(*inds)
result = getter(a)
# 强制转为元组
fixed_result = (result,) if not isinstance(result, tuple) else result
print(fixed_result)  # 输出: ([1, 2, 3],)

方法2：始终用tuple()包裹结果 #

这是我个人最常用的简洁方案，不管inds有几个元素，都用tuple()包装itemgetter的返回值：

from operator import itemgetter

a = [[1, 2, 3], [2, 3, 4], [1, 2], [3, 5]]
inds = [0]

result = tuple(itemgetter(*inds)(a))
print(result)  # 输出: ([1, 2, 3],)

当inds含多个元素时，itemgetter返回的已经是元组，tuple()不会改变它；单个元素时，tuple()会自动把单个元素转换成单元素元组，完美保持格式一致。

方法3：自定义统一格式的getter函数 #

如果经常需要处理这类场景，可以封装一个函数，确保返回格式始终符合预期：

from operator import itemgetter

def consistent_getter(seq, indices):
    getter = itemgetter(*indices)
    result = getter(seq)
    # 处理单个索引的情况
    return (result,) if isinstance(indices, int) or len(indices) == 1 else result

a = [[1, 2, 3], [2, 3, 4], [1, 2], [3, 5]]
print(consistent_getter(a, [0]))  # 输出: ([1, 2, 3],)
print(consistent_getter(a, [0, 2]))  # 输出: ([1, 2, 3], [1, 2])

其实itemgetter的这种行为是为了单元素场景的便捷性，但在批量处理的统一格式需求下就会踩坑，上面的方法都能很好地解决这个问题~

内容的提问来源于stack exchange，提问作者p-value