separate infer_compression method

1 files changed, 11 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