From 9e9bde6c54fec0a3143be2062453611d3ad955e7 Mon Sep 17 00:00:00 2001
From: "Eugeniy E. Mikhailov" <evgmik@gmail.com>
Date: Thu, 20 Jun 2024 22:49:36 -0400
Subject: separate infer_compression method

---
 qolab/data/trace.py          | 26 +++++++++++---------------
 qolab/file_utils/__init__.py | 10 ++++++++++
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index cc51d5b..843d48e 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -1,4 +1,4 @@
-from qolab.file_utils import save_table_with_header
+from qolab.file_utils import save_table_with_header, infer_compression
 import datetime
 import numpy as np
 import yaml
@@ -32,20 +32,16 @@ def loadTraceRawHeaderAndData(fname, tryCompressedIfMissing=True):
                 break
     # we will try to guess if the file compressed
     _open = open
-    for ext in ['gz', 'bz', 'bz2']:
-        b, fext = os.path.splitext(fname)
-        if fext != '.'+ext:
-            continue
-        if fext == '.gz':
-            # TODO improve detection: gzip files have first 2 bytes set to b'\x1f\x8b'
-            import gzip
-            _open = gzip.open
-            break
-        if ( fext == '.bz') or (fext == '.bz2'):
-            # TODO improve detection: bzip files have first 2 bytes set to b'BZ'
-            import bz2
-            _open = bz2.open
-            break
+    compression = infer_compression(fname)
+    print(compression)
+    if compression == 'gzip':
+        # TODO improve detection: gzip files have first 2 bytes set to b'\x1f\x8b'
+        import gzip
+        _open = gzip.open
+    elif compression == 'bzip':
+        # TODO improve detection: bzip files have first 2 bytes set to b'BZ'
+        import bz2
+        _open = bz2.open
     with _open(fname, mode='rb') as tracefile:
         # Reading yaml header prefixed by '% '
         # It sits at the top and below is just data in TSV format
diff --git a/qolab/file_utils/__init__.py b/qolab/file_utils/__init__.py
index e61811e..bcc0824 100644
--- a/qolab/file_utils/__init__.py
+++ b/qolab/file_utils/__init__.py
@@ -69,6 +69,16 @@ 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 infer_compression(fname):
+    """Infers compression algorithm from filename extension"""
+    compression = None # usual suspect
+    b, fext = os.path.splitext(fname)
+    if fext == '.gz':
+        compression = 'gzip'
+    elif ( fext == '.bz') or (fext == '.bz2'):
+        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):
     """Saves output to CSV or TSV file with specially formatted header.
 
-- 
cgit v1.2.3