diff options
| -rw-r--r-- | qolab/file_utils/__init__.py | 14 | ||||
| -rw-r--r-- | tests/test_file_utils.py | 24 |
2 files changed, 36 insertions, 2 deletions
diff --git a/qolab/file_utils/__init__.py b/qolab/file_utils/__init__.py index bcc0824..9b7aeb2 100644 --- a/qolab/file_utils/__init__.py +++ b/qolab/file_utils/__init__.py @@ -79,7 +79,7 @@ def infer_compression(fname): compression = 'bzip' return compression -def save_table_with_header(fname, data, header='', comment_symbol='%', skip_headers_if_file_exist=False, item_format='e', item_separator='\t', compressionmethod=None, compresslevel=9): +def save_table_with_header(fname, data, header='', comment_symbol='%', skip_headers_if_file_exist=False, item_format='e', item_separator='\t', compressionmethod=None, compresslevel=9, match_filename_to_compression=True): """Saves output to CSV or TSV file with specially formatted header. The file is appended if needed. @@ -105,8 +105,20 @@ def save_table_with_header(fname, data, header='', comment_symbol='%', skip_head - gzip : gzip method of compression - bzip : bzip2 method of compression compresslevel : 9 (default) the highest compression, 0 no compression at all, as it is defined for gzip in Lib/gzip.py + match_filename_to_compression: True (default) or False + If True changes the filename suffix in accordance with compression method, + e.g. 'data.dat' -> 'data.dat.gz' if compression is set to 'gzip', + otherwise assumes that users know what they do. """ fname = filename2os_fname(fname) + compression_infered = infer_compression(fname) + if (compression_infered != compressionmethod) and match_filename_to_compression: + if compressionmethod is None: + fname += '.dat' + if compressionmethod == 'gzip': + fname += '.gz' + elif compressionmethod == 'bzip': + fname += '.bz' file_exist_flag = os.path.exists(fname) item_format=str.join('', ['{', f':{item_format}', '}']) _open = open # standard file handler diff --git a/tests/test_file_utils.py b/tests/test_file_utils.py index dd31188..f35999d 100644 --- a/tests/test_file_utils.py +++ b/tests/test_file_utils.py @@ -1,6 +1,7 @@ import pytest +import os -from qolab.file_utils import infer_compression +from qolab.file_utils import infer_compression, save_table_with_header def test_infer_compression(): assert infer_compression('data.dat') is None @@ -9,3 +10,24 @@ def test_infer_compression(): assert infer_compression('data.dat.bz') == 'bzip' assert infer_compression('data.dat.bz2') == 'bzip' +def test_save_table_with_header_match_filename_to_compression(): + import tempfile + data = [ [1, 3], [2, 5] ] + headerstr=['header 1', 'header 2'] + with tempfile.TemporaryDirectory(prefix='qolab_test') as d: + fname = os.path.join(d, 't.dat') + assert save_table_with_header(fname, data, header=headerstr, compressionmethod=None, match_filename_to_compression=True) == (fname) + assert save_table_with_header(fname, data, header=headerstr, compressionmethod='gzip', match_filename_to_compression=True) == (fname+'.gz') + assert save_table_with_header(fname, data, header=headerstr, compressionmethod='bzip', match_filename_to_compression=True) == (fname+'.bz') + + # now cases when requested extension does not match the compression + assert save_table_with_header(fname+'.gz', data, header=headerstr, compressionmethod=None, match_filename_to_compression=True) == (fname+'.gz.dat') + assert save_table_with_header(fname+'.bz', data, header=headerstr, compressionmethod='gzip', match_filename_to_compression=True) == (fname+'.bz.gz') + assert save_table_with_header(fname+'.gz', data, header=headerstr, compressionmethod='bzip', match_filename_to_compression=True) == (fname+'.gz.bz') + + # do as I told cases, extension does not match compression + assert save_table_with_header(fname+'.gz', data, header=headerstr, compressionmethod=None, match_filename_to_compression=False) == (fname+'.gz') + assert save_table_with_header(fname, data, header=headerstr, compressionmethod='gzip', match_filename_to_compression=False) == (fname) + assert save_table_with_header(fname, data, header=headerstr, compressionmethod='bzip', match_filename_to_compression=False) == (fname) + + |