Добавить данные в новый столбец и первую строку в файле CSV

Question

У меня есть функция python, которая создает CSV-файл, используя Postgresql оператор копирования . Мне нужно добавить новый столбец в эту электронную таблицу с именем «UAL» с примерным значением в первой строке, скажем, 30000, но без редактирования оператора копирования. Это текущий код:

copy_sql = 'COPY (
  SELECT
    e.name AS "Employee Name",
    e.title AS "Job Title"
    e.gross AS "Total Pay",
    e.total AS "Total Pay & Benefits",
    e.year AS "Year",
    e.notes AS "Notes",
    j.name AS "Agency",
    e.status AS "Status"

  FROM employee_employee e
  INNER JOIN jurisdiction_jurisdiction j on e.jurisdiction_id = j.id
  WHERE

      e.year = 2011 AND
      j.id = 4479

  ORDER BY "Agency" ASC, "Total Pay & Benefits" DESC
  )'

with open(path, 'w') as csvfile:
    self.cursor.copy_expert(copy_sql, csvfile)

То, что я пытаюсь сделать, это использовать что-то вроде csv.writer для добавления контента вроде этого:

    with open(path, 'w') as csvfile:
        self.cursor.copy_expert(copy_sql, csvfile)
        writer = csv.writer(csvfile)
        writer.writerow('test123')

Но это добавление текста в последнюю строку. Я также не уверен, как добавить новый заголовок столбца. Любой совет?

Answer 1

добавить заголовок легко: напишите заголовок перед вызовом copy_expert.

with open(path, 'w') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["my","super","header"])
    self.cursor.copy_expert(copy_sql, csvfile)

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

Если файлне слишком большой и умещается в памяти, вы можете записать вывод sql в «поддельный» файл:

import io
fakefile = io.StringIO()
self.cursor.copy_expert(copy_sql, fakefile)

теперь перемотайте файл и проанализируйте его как csv, добавьте дополнительный столбец при записиобратно

import csv
fakefile.seek(0)
with open(path, 'w', newline="") as csvfile:
    writer = csv.writer(csvfile)
    reader = csv.reader(fakefile)  # works if copy_expert uses "," as separator, else change it
    writer.writerow(["my","super","header","UAL"])
    for row in reader:
       writer.writerow(row+[30000])

или вместо внутреннего цикла:

    writer.writerows(row+[30000] for row in reader)

А если файл слишком большой, запишите его во временный файл и продолжайте в том же духе (менее производительный)