combine compression code into one · pandas-dev/pandas@74caf72 (original) (raw)
`@@ -287,53 +287,84 @@ def ZipFile(*args, **kwargs):
`
287
287
`ZipFile = zipfile.ZipFile
`
288
288
``
289
289
``
290
``
`-
def _get_handle(path, mode, encoding=None, compression=None, memory_map=False):
`
``
290
`+
def _get_handle(source, mode, encoding=None, compression=None, memory_map=False):
`
291
291
`"""Gets file handle for given path and mode.
`
292
292
` """
`
293
``
`-
if compression is not None:
`
294
``
`-
if encoding is not None and not compat.PY3:
`
``
293
+
``
294
`+
f = source
`
``
295
`+
is_path = isinstance(source, compat.string_types)
`
``
296
+
``
297
`+
in Python 3, convert BytesIO or fileobjects passed with an encoding
`
``
298
`+
if compat.PY3 and isinstance(source, compat.BytesIO):
`
``
299
`+
from io import TextIOWrapper
`
``
300
+
``
301
`+
return TextIOWrapper(source, encoding=encoding)
`
``
302
+
``
303
`+
elif compression is not None:
`
``
304
`+
compression = compression.lower()
`
``
305
`+
if encoding is not None and not compat.PY3 and not is_path:
`
295
306
`msg = 'encoding + compression not yet supported in Python 2'
`
296
307
`raise ValueError(msg)
`
297
308
``
``
309
`+
GZ Compression
`
298
310
`if compression == 'gzip':
`
299
311
`import gzip
`
300
``
`-
f = gzip.GzipFile(path, mode)
`
``
312
+
``
313
`+
f = gzip.GzipFile(source, mode) \
`
``
314
`+
if is_path else gzip.GzipFile(fileobj=source)
`
``
315
+
``
316
`+
BZ Compression
`
301
317
`elif compression == 'bz2':
`
302
318
`import bz2
`
303
``
`-
f = bz2.BZ2File(path, mode)
`
``
319
+
``
320
`+
if is_path:
`
``
321
`+
f = bz2.BZ2File(source, mode)
`
``
322
+
``
323
`+
else:
`
``
324
`+
f = bz2.BZ2File(source) if compat.PY3 else StringIO(
`
``
325
`+
bz2.decompress(source.read()))
`
``
326
`+
Python 2's bz2 module can't take file objects, so have to
`
``
327
`+
run through decompress manually
`
``
328
+
``
329
`+
ZIP Compression
`
304
330
`elif compression == 'zip':
`
305
331
`import zipfile
`
306
``
`-
zip_file = zipfile.ZipFile(path)
`
``
332
`+
zip_file = zipfile.ZipFile(source)
`
307
333
`zip_names = zip_file.namelist()
`
308
334
``
309
335
`if len(zip_names) == 1:
`
310
``
`-
file_name = zip_names.pop()
`
311
``
`-
f = zip_file.open(file_name)
`
``
336
`+
f = zip_file.open(zip_names.pop())
`
312
337
`elif len(zip_names) == 0:
`
313
338
`raise ValueError('Zero files found in ZIP file {}'
`
314
``
`-
.format(path))
`
``
339
`+
.format(source))
`
315
340
`else:
`
316
341
`raise ValueError('Multiple files found in ZIP file.'
`
317
342
`' Only one file per ZIP :{}'
`
318
343
` .format(zip_names))
`
``
344
+
``
345
`+
XZ Compression
`
319
346
`elif compression == 'xz':
`
320
347
`lzma = compat.import_lzma()
`
321
``
`-
f = lzma.LZMAFile(path, mode)
`
``
348
`+
f = lzma.LZMAFile(source, mode)
`
``
349
+
322
350
`else:
`
323
``
`-
raise ValueError('Unrecognized compression type: %s' %
`
324
``
`-
compression)
`
``
351
`+
raise ValueError('Unrecognized compression: %s' % compression)
`
``
352
+
325
353
`if compat.PY3:
`
326
354
`from io import TextIOWrapper
`
``
355
+
327
356
`f = TextIOWrapper(f, encoding=encoding)
`
``
357
+
328
358
`return f
`
329
``
`-
else:
`
``
359
+
``
360
`+
elif is_path:
`
330
361
`if compat.PY3:
`
331
362
`if encoding:
`
332
``
`-
f = open(path, mode, encoding=encoding)
`
``
363
`+
f = open(source, mode, encoding=encoding)
`
333
364
`else:
`
334
``
`-
f = open(path, mode, errors='replace')
`
``
365
`+
f = open(source, mode, errors='replace')
`
335
366
`else:
`
336
``
`-
f = open(path, mode)
`
``
367
`+
f = open(source, mode)
`
337
368
``
338
369
`if memory_map and hasattr(f, 'fileno'):
`
339
370
`try:
`