33"""
44
55import csv as csvlib
6- from io import StringIO
6+ from io import StringIO , TextIOWrapper
77import os
88from typing import Hashable , List , Mapping , Optional , Sequence , Union
99import warnings
10- from zipfile import ZipFile
1110
1211import numpy as np
1312
@@ -159,38 +158,29 @@ def save(self) -> None:
159158 """
160159 Create the writer & save.
161160 """
162- # GH21227 internal compression is not used when file-like passed.
163- if self .compression and hasattr (self .path_or_buf , "write" ):
161+ # GH21227 internal compression is not used for non-binary handles.
162+ if (
163+ self .compression
164+ and hasattr (self .path_or_buf , "write" )
165+ and "b" not in self .mode
166+ ):
164167 warnings .warn (
165- "compression has no effect when passing file-like object as input." ,
168+ "compression has no effect when passing a non-binary object as input." ,
166169 RuntimeWarning ,
167170 stacklevel = 2 ,
168171 )
169-
170- # when zip compression is called.
171- is_zip = isinstance (self .path_or_buf , ZipFile ) or (
172- not hasattr (self .path_or_buf , "write" ) and self .compression == "zip"
172+ self .compression = None
173+
174+ # get a handle or wrap an existing handle to take care of 1) compression and
175+ # 2) text -> byte conversion
176+ f , handles = get_handle (
177+ self .path_or_buf ,
178+ self .mode ,
179+ encoding = self .encoding ,
180+ errors = self .errors ,
181+ compression = dict (self .compression_args , method = self .compression ),
173182 )
174183
175- if is_zip :
176- # zipfile doesn't support writing string to archive. uses string
177- # buffer to receive csv writing and dump into zip compression
178- # file handle. GH21241, GH21118
179- f = StringIO ()
180- close = False
181- elif hasattr (self .path_or_buf , "write" ):
182- f = self .path_or_buf
183- close = False
184- else :
185- f , handles = get_handle (
186- self .path_or_buf ,
187- self .mode ,
188- encoding = self .encoding ,
189- errors = self .errors ,
190- compression = dict (self .compression_args , method = self .compression ),
191- )
192- close = True
193-
194184 try :
195185 # Note: self.encoding is irrelevant here
196186 self .writer = csvlib .writer (
@@ -206,29 +196,23 @@ def save(self) -> None:
206196 self ._save ()
207197
208198 finally :
209- if is_zip :
210- # GH17778 handles zip compression separately.
211- buf = f .getvalue ()
212- if hasattr (self .path_or_buf , "write" ):
213- self .path_or_buf .write (buf )
214- else :
215- compression = dict (self .compression_args , method = self .compression )
216-
217- f , handles = get_handle (
218- self .path_or_buf ,
219- self .mode ,
220- encoding = self .encoding ,
221- errors = self .errors ,
222- compression = compression ,
223- )
224- f .write (buf )
225- close = True
226- if close :
199+ if self .should_close :
227200 f .close ()
228- for _fh in handles :
229- _fh .close ()
230- elif self .should_close :
201+ elif (
202+ isinstance (f , TextIOWrapper )
203+ and not f .closed
204+ and f != self .path_or_buf
205+ and hasattr (self .path_or_buf , "write" )
206+ ):
207+ # get_handle uses TextIOWrapper for non-binary handles. TextIOWrapper
208+ # closes the wrapped handle if it is not detached.
209+ f .flush () # make sure everything is written
210+ f .detach () # makes f unusable
211+ del f
212+ elif f != self .path_or_buf :
231213 f .close ()
214+ for _fh in handles :
215+ _fh .close ()
232216
233217 def _save_header (self ):
234218 writer = self .writer
0 commit comments