move decompress_file to util/testing.py

Author: reidy-p

pandas/tests/series/test_io.py

@@ -139,24 +139,6 @@ def test_to_csv_path_is_none(self):
  csv_str = s.to_csv(path=None)
  assert isinstance(csv_str, str)
 
- def decompress_file(self, src_path, compression):
- if compression is None:
- f = open(src_path, 'rb')
- elif compression == 'gzip':
- import gzip
- f = gzip.open(src_path, 'rb')
- elif compression == 'bz2':
- import bz2
- f = bz2.BZ2File(src_path, 'rb')
- elif compression == 'xz':
- lzma = compat.import_lzma()
- f = lzma.LZMAFile(src_path, 'rb')
- else:
- msg = 'Unrecognized compression type: {}'.format(compression)
- raise ValueError(msg)
-
- return f
-
  @pytest.mark.parametrize('compression', [
  None,
  'gzip',
@@ -178,12 +160,12 @@ def test_to_csv_compression(self, compression):
  assert_series_equal(s, rs)
 
  # explicitly ensure file was compressed
- f = self.decompress_file(filename, compression=compression)
+ f = tm.decompress_file(filename, compression=compression)
  text = f.read().decode('utf8')
  assert s.name in text
  f.close()
 
- f = self.decompress_file(filename, compression=compression)
+ f = tm.decompress_file(filename, compression=compression)
  assert_series_equal(s, pd.read_csv(f, index_col=0, squeeze=True))
  f.close()
 

pandas/util/testing.py

@@ -161,6 +161,39 @@ def round_trip_localpath(writer, reader, path=None):
  obj = reader(LocalPath(path))
  return obj
 
+def decompress_file(path, compression):
+ """
+ Open a compressed file and return a file object
+
+ Parameters
+ ----------
+ path : str
+ The path where the file is read from
+
+ compression : {'gzip', 'bz2', 'xz', None}
+ Name of the decompression to use
+
+ Returns
+ -------
+ f : file object
+ """
+
+ if compression is None:
+ f = open(path, 'rb')
+ elif compression == 'gzip':
+ import gzip
+ f = gzip.open(path, 'rb')
+ elif compression == 'bz2':
+ import bz2
+ f = bz2.BZ2File(path, 'rb')
+ elif compression == 'xz':
+ lzma = compat.import_lzma()
+ f = lzma.LZMAFile(path, 'rb')
+ else:
+ msg = 'Unrecognized compression type: {}'.format(compression)
+ raise ValueError(msg)
+
+ return f
 
 def assert_almost_equal(left, right, check_exact=False,
  check_dtype='equiv', check_less_precise=False,