diff options
| -rw-r--r-- | qolab/file_utils/__init__.py | 18 |
1 files changed, 13 insertions, 5 deletions
diff --git a/qolab/file_utils/__init__.py b/qolab/file_utils/__init__.py index 45bcd4b..456fc2f 100644 --- a/qolab/file_utils/__init__.py +++ b/qolab/file_utils/__init__.py @@ -69,20 +69,28 @@ def get_next_data_file(prefix, savepath, run_number=None, datestr=None, date_for fname = os.path.join(savepath, f'{prefix}_{datestr}_{run_number:05d}.{extension}') return(fname) -def save_table_with_header(fname, data, header='', comment_symbol='%', skip_headers_if_file_exist=False, item_format='e', item_separator='\t'): +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): # itemFormat examples: 'e', '.15e', 'f' + # `compressionmethod` can be + # None - no compression + # "gzip" - gzip method of compression + # `compresslevel`: 9 the highest compression, 0 no compression at all, as it is defined for gzip in Lib/gzip.py fname = filename2os_fname(fname) file_exist_flag = os.path.exists(fname) item_format=str.join('', ['{', f':{item_format}', '}']) - with open(fname, 'ab') as f: + _open = open # standard file handler + if compressionmethod == 'gzip': + import gzip + _open = lambda fname, mode: gzip.open( fname, mode, compresslevel = compresslevel) + with _open(fname, 'ab') as f: if not (file_exist_flag and skip_headers_if_file_exist): for l in header: - f.write(f'{comment_symbol} {l}\n') + f.write(f'{comment_symbol} {l}\n'.encode('utf-8')) if data is not None: for r in data: l=item_separator.join( map(item_format.format, r)) - f.write(l) - f.write('\n') + f.write(l.encode('utf-8')) + f.write('\n'.encode('utf-8')) f.close() return(fname) |