added gzip compression for file saving

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)