diff --git a/pandas/tests/io/parser/compression.py b/pandas/tests/io/parser/compression.py index 4291d59123e8b..01c6620e50d37 100644 --- a/pandas/tests/io/parser/compression.py +++ b/pandas/tests/io/parser/compression.py @@ -12,6 +12,13 @@ import pandas.util.testing as tm import pandas.util._test_decorators as td +import gzip +import bz2 +try: + lzma = compat.import_lzma() +except ImportError: + lzma = None + class CompressionTests(object): @@ -64,83 +71,36 @@ def test_zip(self): pytest.raises(zipfile.BadZipfile, self.read_csv, f, compression='zip') - def test_gzip(self): - import gzip - - with open(self.csv1, 'rb') as data_file: - data = data_file.read() - expected = self.read_csv(self.csv1) - - with tm.ensure_clean() as path: - tmp = gzip.GzipFile(path, mode='wb') - tmp.write(data) - tmp.close() - - result = self.read_csv(path, compression='gzip') - tm.assert_frame_equal(result, expected) - - with open(path, 'rb') as f: - result = self.read_csv(f, compression='gzip') - tm.assert_frame_equal(result, expected) - - with tm.ensure_clean('test.gz') as path: - tmp = gzip.GzipFile(path, mode='wb') - tmp.write(data) - tmp.close() - result = self.read_csv(path, compression='infer') - tm.assert_frame_equal(result, expected) - - def test_bz2(self): - import bz2 + @pytest.mark.parametrize('compress_type, compress_method, ext', [ + ('gzip', gzip.GzipFile, 'gz'), + ('bz2', bz2.BZ2File, 'bz2'), + pytest.param('xz', getattr(lzma, 'LZMAFile', None), 'xz', + marks=td.skip_if_no_lzma) + ]) + def test_other_compression(self, compress_type, compress_method, ext): with open(self.csv1, 'rb') as data_file: data = data_file.read() expected = self.read_csv(self.csv1) with tm.ensure_clean() as path: - tmp = bz2.BZ2File(path, mode='wb') + tmp = compress_method(path, mode='wb') tmp.write(data) tmp.close() - result = self.read_csv(path, compression='bz2') + result = self.read_csv(path, compression=compress_type) tm.assert_frame_equal(result, expected) - pytest.raises(ValueError, self.read_csv, - path, compression='bz3') + if compress_type == 'bz2': + pytest.raises(ValueError, self.read_csv, + path, compression='bz3') with open(path, 'rb') as fin: - result = self.read_csv(fin, compression='bz2') - tm.assert_frame_equal(result, expected) - - with tm.ensure_clean('test.bz2') as path: - tmp = bz2.BZ2File(path, mode='wb') - tmp.write(data) - tmp.close() - result = self.read_csv(path, compression='infer') - tm.assert_frame_equal(result, expected) - - @td.skip_if_no_lzma - def test_xz(self): - lzma = compat.import_lzma() - - with open(self.csv1, 'rb') as data_file: - data = data_file.read() - expected = self.read_csv(self.csv1) - - with tm.ensure_clean() as path: - tmp = lzma.LZMAFile(path, mode='wb') - tmp.write(data) - tmp.close() - - result = self.read_csv(path, compression='xz') - tm.assert_frame_equal(result, expected) - - with open(path, 'rb') as f: - result = self.read_csv(f, compression='xz') + result = self.read_csv(fin, compression=compress_type) tm.assert_frame_equal(result, expected) - with tm.ensure_clean('test.xz') as path: - tmp = lzma.LZMAFile(path, mode='wb') + with tm.ensure_clean('test.{}'.format(ext)) as path: + tmp = compress_method(path, mode='wb') tmp.write(data) tmp.close() result = self.read_csv(path, compression='infer')