Python中如何遍历可为单个对象或对象列表/迭代器的变量？

问题 #

我有一个Python函数，参数支持int类型或int列表，想要对参数进行遍历操作：

from typing import List, Union


def myfunc(arg: Union[int, List[int]]) -> List[int]:
    x = []
    for value in arg:
        x.append(2 * value)

    return x

显然，传入列表时函数能正常运行，但传入int类型参数会抛出TypeError——因为int不是可迭代对象。

我通常会这样处理：

def myfunc(arg: Union[int, List[int]]) -> List[int]:
    x = []
    arglist = arg if isinstance(arg, list) else [arg]
    for value in arglist:
        x.append(2 * value)

    return x

也就是如果参数不是列表，就把它放进一个列表里再遍历。

我也可以自己写个函数，把非可迭代对象转换成可迭代对象：

def make_iterator(x):
    try:
        for v in x:
            yield v
    except TypeError:
        yield x


def myfunc(arg: Union[int, List[int]]) -> List[int]:
    x = []
    for value in make_iterator(arg):
        x.append(2 * value)

    return x

但我想知道，Python有没有内置函数或者标准库模块可以直接实现这个功能？

解决方案 #

Python标准库中没有专门提供这类“将非可迭代对象转为单元素可迭代对象”的内置函数，但可以通过现有工具或者更简洁的逻辑实现类似效果：

1. 扩展可迭代对象判断范围 #

如果你想支持更多可迭代类型（而非仅列表），可以用collections.abc.Iterable来判断，同时注意排除字符串（避免单个字符串被拆成字符遍历）：

from typing import List, Union
from collections.abc import Iterable

def myfunc(arg: Union[int, List[int]]) -> List[int]:
    # 单独处理字符串，避免拆分字符
    if isinstance(arg, str):
        target = [arg]
    else:
        target = arg if isinstance(arg, Iterable) else [arg]
    return [2 * val for val in target]

2. 借助itertools.chain实现 #

利用itertools.chain拼接可迭代对象的特性，结合异常判断来处理：

from typing import List, Union
import itertools

def myfunc(arg: Union[int, List[int]]) -> List[int]:
    try:
        # 尝试生成迭代器，判断是否为可迭代对象
        iter(arg)
        # 排除字符串
        target = arg if not isinstance(arg, str) else [arg]
    except TypeError:
        target = [arg]
    return [2 * val for val in itertools.chain(target)]

这个逻辑和你自己写的make_iterator类似，只是借助了标准库工具。

3. 针对原场景的极简写法 #

如果只需要支持int和List[int]的场景，你最初的写法已经足够简洁，完全不需要依赖额外库：

from typing import List, Union

def myfunc(arg: Union[int, List[int]]) -> List[int]:
    return [2 * val for val in (arg if isinstance(arg, list) else [arg])]

总的来说，这类需求比较场景化，Python没有专门的内置函数来处理，但自己实现简单判断或者用上述方式处理就足够高效。

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