Добавить столбец (со значениями) во время записи в csv с помощью csv.DictWriter в python

Question

У меня есть файл yml следующим образом:

Stocks:
 - name: shop1
   apples: 100
   mangoes: 200

 - name: shop2
   apples: 20
   mangoes: 50
   bananas: 100
   garbage: 'something'

Я хочу создать файл данных TSV (или CSV или PANDAS), который выглядит следующим образом:

name    category  apples   mangoes   bananas
shop1   fruits    100      200       0
shop2   fruits    20       50        100

Мой код на данный момент выглядит следующим образом:

import yaml, csv

csv_columns = ['name','apples','mangoes','bananas', 'category']

with open('fruits.yml', 'r') as f:
    doc = yaml.load(f)


csv_file = "somename.csv"
try:
    with open(csv_file, 'w') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=csv_columns, extrasaction='ignore')
        writer.writeheader()
   for item in doc["Stocks"]:
        item.update( {"category":"fruits"})
   for data in doc["Stocks"]:
            writer.writerow(data)
except IOError:
    print("I/O error")

Работает нормально, за исключением того, что Я не знаю правильного или элегантного способа добавить столбец category с тем же значением, повторяемым в строках. Это значение фактически является именем файла yaml.

Answer 1

Вот что вам нужно, в repl.it: https://repl.it/repls/MellowDisfiguredGroupware

Позвольте мне объяснить:

Сначала нужно изменить csv_columns

csv_columns = ['name','category','apples','mangoes','bananas']

Во-вторых, исходя из заданного вами yaml, в словаре doc нет ключа websites, только ключ Stocks. Я соответственно отредактировал письменную часть.

В-третьих, вам нужно добавить ключ category со значением fruits:

    try:
    with open(csv_file, 'w') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=csv_columns, extrasaction='ignore')
        writer.writeheader()
        for data in doc["Stocks"]:
            data['category'] = 'fruits'
            writer.writerow(data)
except IOError:
    print("I/O error")

Answer 2

При записи строки данных вы можете использовать понимание списка следующим образом.

for data in doc["websites"]:
    writer.writerow([data[0], "fruits"]+data[1:])