Skip to content

Commit 9567900

Browse files
committed
regression to 0.22 and add zipfile support
1 parent d3dac82 commit 9567900

File tree

1 file changed

+43
-19
lines changed

1 file changed

+43
-19
lines changed

pandas/io/formats/csvs.py

Lines changed: 43 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
import numpy as np
1010

1111
from pandas.core.dtypes.missing import notna
12-
from pandas.core.dtypes.inference import is_file_like
1312
from pandas.core.index import Index, MultiIndex
1413
from pandas import compat
1514
from pandas.compat import (StringIO, range, zip)
@@ -128,19 +127,26 @@ def save(self):
128127
else:
129128
encoding = self.encoding
130129

131-
# PR 21300 uses string buffer to receive csv writing and dump into
132-
# file-like output with compression as option. GH 21241, 21118
133-
f = StringIO()
134-
if not is_file_like(self.path_or_buf):
135-
# path_or_buf is path
136-
path_or_buf = self.path_or_buf
137-
elif hasattr(self.path_or_buf, 'name'):
138-
# path_or_buf is file handle
139-
path_or_buf = self.path_or_buf.name
130+
if (not hasattr(self.path_or_buf, 'write') and
131+
self.compression == 'zip'):
132+
is_zip = True
140133
else:
141-
# path_or_buf is file-like IO objects.
134+
is_zip = False
135+
136+
if is_zip:
137+
# zipfile doesn't support writing string to archive. uses string
138+
# buffer to receive csv writing and dump into zip compression
139+
# file handle. GH 21241, 21118
140+
f = StringIO()
141+
close = False
142+
elif hasattr(self.path_or_buf, 'write'):
142143
f = self.path_or_buf
143-
path_or_buf = None
144+
close = False
145+
else:
146+
f, handles = _get_handle(self.path_or_buf, self.mode,
147+
encoding=encoding,
148+
compression=self.compression)
149+
close = True
144150

145151
try:
146152
writer_kwargs = dict(lineterminator=self.line_terminator,
@@ -157,17 +163,35 @@ def save(self):
157163
self._save()
158164

159165
finally:
160-
# GH 17778 handles zip compression for byte strings separately.
161-
buf = f.getvalue()
162-
if path_or_buf:
163-
f, handles = _get_handle(path_or_buf, self.mode,
164-
encoding=encoding,
165-
compression=self.compression)
166-
f.write(buf)
166+
if is_zip:
167+
# GH 17778 handles zip compression separately.
168+
buf = f.getvalue()
169+
try:
170+
self.path_or_buf.write(buf)
171+
except AttributeError:
172+
f, handles = _get_handle(self.path_or_buf, self.mode,
173+
encoding=encoding,
174+
compression=self.compression)
175+
f.write(buf)
176+
close = True
177+
if close:
167178
f.close()
168179
for _fh in handles:
169180
_fh.close()
170181

182+
# GH 17778 handles zip compression for byte strings separately.
183+
# if path is not None:
184+
# buf = f.getvalue()
185+
# f, handles = _get_handle(path, self.mode,
186+
# encoding=encoding,
187+
# compression=self.compression)
188+
# f.write(buf)
189+
#
190+
# if not py2zip:
191+
# f.close()
192+
# for _fh in handles:
193+
# _fh.close()
194+
171195
def _save_header(self):
172196

173197
writer = self.writer

0 commit comments

Comments
 (0)