11
11
from pandas .compat import range , lmap
12
12
import pandas .core .common as com
13
13
from pandas .core import ops
14
- from pandas .io .common import _get_handle
14
+ from pandas .io .common import (
15
+ _compression_to_extension ,
16
+ _get_handle ,
17
+ )
15
18
import pandas .util .testing as tm
16
19
17
20
@@ -217,13 +220,23 @@ def test_standardize_mapping():
217
220
Series (100 * [0.123456 , 0.234567 , 0.567567 ], name = 'X' )])
218
221
@pytest .mark .parametrize ('method' , ['to_pickle' , 'to_json' , 'to_csv' ])
219
222
def test_compression_size (obj , method , compression_only ):
220
-
221
- with tm .ensure_clean () as filename :
222
- getattr (obj , method )(filename , compression = compression_only )
223
- compressed = os .path .getsize (filename )
224
- getattr (obj , method )(filename , compression = None )
225
- uncompressed = os .path .getsize (filename )
226
- assert uncompressed > compressed
223
+ """
224
+ Tests that compression is occurring by comparing to the bytes on disk of
225
+ the uncompressed file.
226
+ """
227
+ extension = _compression_to_extension [compression_only ]
228
+ to_method = getattr (obj , method )
229
+ with tm .ensure_clean ('no-compression' ) as path :
230
+ to_method (path , compression = None )
231
+ no_compression_size = os .path .getsize (path )
232
+ with tm .ensure_clean ('explicit-compression' + extension ) as path :
233
+ to_method (path , compression = compression_only )
234
+ explicit_compression_size = os .path .getsize (path )
235
+ with tm .ensure_clean ('inferred-compression' + extension ) as path :
236
+ to_method (path ) # assumes that compression='infer' is the default
237
+ inferred_compression_size = os .path .getsize (path )
238
+ assert (no_compression_size > explicit_compression_size ==
239
+ inferred_compression_size )
227
240
228
241
229
242
@pytest .mark .parametrize ('obj' , [
@@ -233,7 +246,6 @@ def test_compression_size(obj, method, compression_only):
233
246
Series (100 * [0.123456 , 0.234567 , 0.567567 ], name = 'X' )])
234
247
@pytest .mark .parametrize ('method' , ['to_csv' , 'to_json' ])
235
248
def test_compression_size_fh (obj , method , compression_only ):
236
-
237
249
with tm .ensure_clean () as filename :
238
250
f , _handles = _get_handle (filename , 'w' , compression = compression_only )
239
251
with f :
0 commit comments