Skip to content

Commit d9ac263

Browse files
committed
make decompress_file a context manager and create compression fixture
1 parent e299354 commit d9ac263

File tree

3 files changed

+21
-15
lines changed

3 files changed

+21
-15
lines changed

pandas/tests/conftest.py

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
import pytest
2+
import pandas
3+
import pandas.util._test_decorators as td
4+
5+
6+
@pytest.fixture(params=[None, 'gzip', 'bz2',
7+
pytest.param('xz', marks=td.skip_if_no_lzma)])
8+
def compression(request):
9+
"""
10+
Fixture for trying common compression types in compression tests
11+
"""
12+
return request.param

pandas/tests/frame/test_to_csv.py

+7-14
Original file line numberDiff line numberDiff line change
@@ -920,12 +920,6 @@ def test_to_csv_path_is_none(self):
920920
recons = pd.read_csv(StringIO(csv_str), index_col=0)
921921
assert_frame_equal(self.frame, recons)
922922

923-
@pytest.mark.parametrize('compression', [
924-
None,
925-
'gzip',
926-
'bz2',
927-
pytest.param('xz', marks=td.skip_if_no_lzma),
928-
])
929923
def test_to_csv_compression(self, compression):
930924

931925
df = DataFrame([[0.123456, 0.234567, 0.567567],
@@ -941,14 +935,13 @@ def test_to_csv_compression(self, compression):
941935
assert_frame_equal(df, rs)
942936

943937
# explicitly make sure file is compressed
944-
f = tm.decompress_file(filename, compression)
945-
text = f.read().decode('utf8')
946-
for col in df.columns:
947-
assert col in text
948-
f.close()
949-
950-
f = tm.decompress_file(filename, compression)
951-
assert_frame_equal(df, read_csv(f, index_col=0))
938+
with tm.decompress_file(filename, compression) as fh:
939+
text = fh.read().decode('utf8')
940+
for col in df.columns:
941+
assert col in text
942+
943+
with tm.decompress_file(filename, compression) as fh:
944+
assert_frame_equal(df, read_csv(fh, index_col=0))
952945

953946
def test_to_csv_compression_value_error(self):
954947
# GH7615

pandas/util/testing.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ def round_trip_localpath(writer, reader, path=None):
162162
return obj
163163

164164

165+
@contextmanager
165166
def decompress_file(path, compression):
166167
"""
167168
Open a compressed file and return a file object
@@ -194,7 +195,7 @@ def decompress_file(path, compression):
194195
msg = 'Unrecognized compression type: {}'.format(compression)
195196
raise ValueError(msg)
196197

197-
return f
198+
yield f
198199

199200

200201
def assert_almost_equal(left, right, check_exact=False,

0 commit comments

Comments
 (0)