Как определить два способа итерации в зависимости от ввода

Question

Какая наиболее питоническая форма для определения двух способов итерации. Например, у меня есть этот оригинальный код:

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

Есть ли способ условно изменить (a, b), c на a, b, c в операторе for

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

и простоПользователь должен знать способы возврата по этому второму аргументу.