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

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):
    # 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 2

通过一个可以正确评估条件的函数：

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 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

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