added automatic check if compressed trace file exist if the main is missing

1 files changed, 10 insertions, 3 deletions

diff --git a/qolab/data/trace.py b/qolab/data/trace.py
index 7a139b8..674922a 100644
--- a/qolab/data/trace.py
+++ b/qolab/data/trace.py
@@ -3,6 +3,7 @@ import datetime
 import numpy as np
 import yaml
 import pandas
+import os
 
 def headerFromDictionary(d, prefix=''):
     header = []
@@ -16,9 +17,15 @@ def from_timestamps_to_dates(timestamps):
     dates = [datetime.datetime.fromtimestamp(float(ts)) for ts in timestamps]
     return(dates)
 
-def loadTraceRawHeaderAndData(fname):
+def loadTraceRawHeaderAndData(fname, tryCompressedIfMissing=True):
+    # Attempts to load a compressed file if the main is missing
+    # and `tryCompressedIfMissing` is set to `True`.
+    # E.g. if we try to load 'data.dat' file and it is missing, attempt to load 'data.dat.gz' 
+    # 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'
     # we will try to guess if the file compressed
     _open = open
     if fname[-3:] == '.gz':
@@ -44,8 +51,8 @@ def loadTraceRawHeaderAndData(fname):
         data = df.to_numpy()
     return(header, data)
 
-def loadTrace(fname):
-    (header, data) = loadTraceRawHeaderAndData(fname)
+def loadTrace(fname, tryCompressedIfMissing=True):
+    (header, data) = loadTraceRawHeaderAndData(fname, tryCompressedIfMissing=tryCompressedIfMissing)
     return traceFromHeaderAndData(header, data)
 
 def traceFromHeaderAndData(header, data=None):