added bzip compression reading for traces

1 files changed, 21 insertions, 6 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 674922a..bda4da5 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -24,14 +24,29 @@ def loadTraceRawHeaderAndData(fname, tryCompressedIfMissing=True):
     # this option is mainly for compatibility with old scripts which are not aware of compressed files
     headerstr=[]
     data = None
-    if (not os.path.exists(fname)) and tryCompressedIfMissing and os.path.exists(fname+'.gz'):
-        fname += '.gz'
+    if (not os.path.exists(fname)) and tryCompressedIfMissing:
+        if os.path.exists(fname+'.gz'):
+            fname += '.gz'
+        if os.path.exists(fname+'.bz'):
+            fname += '.bz'
+        if os.path.exists(fname+'.bz2'):
+            fname += '.bz2'
     # we will try to guess if the file compressed
     _open = open
-    if fname[-3:] == '.gz':
-        # TODO improve detection: gzip files have first 2 bytes set to b'\x1f\x8b'
-        import gzip
-        _open = gzip.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
     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