Панды to_csv в GzipFile в Python 3 не работают

Question

Сохранение фрейма данных Pandas в gzipped csv в памяти работает следующим образом в Python 2.7 (Pandas 0.22.0):

from io import BytesIO
import gzip
import pandas as pd
df = pd.DataFrame.from_dict({'a': ['a', 'b', 'c']})
s = BytesIO()
f = gzip.GzipFile(fileobj=s, mode='wb', filename='file.csv')
df.to_csv(f)
s.seek(0)
content = s.getvalue()

Однако в Python 3.6 (Pandas 0.22.0) тот же код выдает ошибку при вызове to_csv:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "lib/python3.6/site-packages/pandas/core/frame.py", line 1524, in to_csv
    formatter.save()
  File "lib/python3.6/site-packages/pandas/io/formats/format.py", line 1652, in save
    self._save()
  File "lib/python3.6/site-packages/pandas/io/formats/format.py", line 1740, in _save
    self._save_header()
  File "lib/python3.6/site-packages/pandas/io/formats/format.py", line 1708, in _save_header
    writer.writerow(encoded_labels)
  File "miniconda3/lib/python3.6/gzip.py", line 260, in write
    data = memoryview(data)
TypeError: memoryview: a bytes-like object is required, not 'str'

Как мне решить эту проблему? Нужно ли как-то изменить объект GzipFile, чтобы to_csv обрабатывал его правильно?

Чтобы уточнить, я хочу создать файл gzipped в памяти (переменная content), чтобы я мог сохранить его в Amazon S3, используя Boto 3 put_object позже.

Answer 1

Вы можете использовать StringIO:

from io import StringIO
buf = StringIO()
df.to_csv(buf)
f = gzip.GzipFile(fileobj=s, mode='wb', filename='file.csv')
f.write(buf.getvalue().encode())
f.flush()

Обратите внимание также на добавленное f.flush() - согласно моему опыту без этой строки GzipFile может в некоторых случаях случайным образом не сбрасывать данные, что приводит к повреждению архива.

Или как полный пример, основанный на вашем коде:

from io import BytesIO
import gzip
import pandas as pd
from io import StringIO
df = pd.DataFrame.from_dict({'a': ['a', 'b', 'c']})
s = BytesIO()
buf = StringIO()
f = gzip.GzipFile(fileobj=s, mode='wb', filename='file.csv')
df.to_csv(buf)
f.write(buf.getvalue().encode())
f.flush()
s.seek(0)
content = s.getvalue()