Попытка написать Python CSV экстрактор

Question

Я новичок в программировании, и это моя первая настоящая программа, которую я пытаюсь написать.

Итак, у меня есть этот огромный CSV-файл (сотни столбцов и тысячи строк), где я пытаюсь извлечь только несколько столбцов на основе значения в поле.Это работает нормально, и я получаю хороший вывод, но проблема возникает, когда я пытаюсь инкапсулировать ту же логику в функцию.он возвращает только первую извлеченную строку, однако печать работает нормально.

Я часами играл в это и читал другие примеры здесь, и теперь мой ум месиво.

import csv
import sys

newlogfile = csv.reader(open(sys.argv[1], 'rb'))
outLog = csv.writer(open('extracted.csv', 'w'))

def rowExtractor(logfile):
    for row in logfile:
        if row[32] == 'No':
            a = []
            a.append(row[44])
            a.append(row[58])
            a.append(row[83])
            a.append(row[32])
            return a

outLog.writerow(rowExtractor(newlogfile))

Answer 1

В вашей функции есть оператор return ... когда он попадает в эту строку, он возвращается (таким образом, завершая ваш цикл). Вам нужно yield вместо.

См. Что делает ключевое слово yield в Python?

Answer 2

Вы выходите преждевременно.Когда вы помещаете return a в цикл for, return вызывается на первой итерации.Это означает, что выполняется только первая итерация.

Простой способ сделать это - сделать:

def rowExtractor(logfile):
    #output holds all of the rows
    ouput = []
    for row in logfile:
        if row[32] == 'No':
            a = []
            a.append(row[44])
            a.append(row[58])
            a.append(row[83])
            a.append(row[32])
            output.append(a)
    #notice that the return statement is outside of the for-loop
    return output
outLog.writerows(rowExtractor(newlogfile))

Вы также можете рассмотреть возможность использования yield