如何根据输入定义两种迭代方式

Question

定义两种迭代方式的最 Pythonic 形式是什么。例如，我有这个原始代码：

def f1(cat_gen):
    for (a, b), c in cat_gen:
        if some condition:
            yield (a, b), c

但是，取决于cat_gen，我需要以这种方式迭代：

def f1(cat_gen):
    for a, b, c in cat_gen:
        if some condition:
            yield a, b, c

有没有办法在for 语句中将(a, b), c 有条件地更改为a, b, c

Answer 1

传递一个可以正确评估条件的函数：

def f1(cat_gen, pred):
    for item in cat_gen:
        if pred(item):
            yield item

f1(flat, lambda a, b, c: ...)
f1(nested, lambda ab, c: ...)

或者，在将可迭代对象传递给 f1 之前将嵌套元组展平：

def f1(cat_gen):
    for a, b, c in cat_gen:
        if ...:
            yield a, b, c

f1(map(lambda ab, c: (ab[0], ab[1], c), nested))

Answer 2

你可以这样定义

def f1(cat_gen):
    # Figure out how to iterate, store it in condition_to_iterate
    for item in cat_gen:
        if condition_to_iterate:
            (a, b), c = item
        else:
            a, b, c = item
        # Do whatever you need with a, b, c

Answer 3

如果您想将其保留为 1 个函数，并使用两种不同的返回形式（可能不是最干净的决定，但取决于实现），您需要执行以下操作：

def f1(cat_gen, yield_method = 0):
    for a, b, c in cat_gen:
        if some condition:
            if yield_method:
                yield a, b, c
            else:
                yield (a, b), c

并让用户知道第二个参数的返回模式。