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

3 files changed, 27 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):
diff --git a/tests/test_trace.py b/tests/test_trace.py
index 981d1b7..3d77df3 100644
--- a/tests/test_trace.py
+++ b/tests/test_trace.py
@@ -1,3 +1,4 @@
+import pytest
 from qolab.data.trace import loadTrace
 import numpy as np
 
@@ -17,3 +18,19 @@ def test_load_gzip_compressed_v0dot1_trace():
     data = tr.getData()
     assert np.all( (data - np.array([[1], [3], [2], [5]])) == 0 )
 
+def test_tryCompressedIfMissing():
+    fname = 'tests/trace_test_data/only_compressed_file1.dat'
+    # first we check that the guess is working
+    tr = loadTrace(fname, tryCompressedIfMissing=True)
+    cfg = tr.getConfig()
+    assert cfg['config']['version'] == '0.1'
+    assert cfg['config']['model'] == 'Trace'
+    data = tr.getData()
+    assert np.all( (data - np.array([[1], [3], [2], [5]])) == 0 )
+
+    # now we disable search for compressed version
+    with pytest.raises(FileNotFoundError) as exc_info:
+        tr = loadTrace(fname, tryCompressedIfMissing=False)
+
+
+
diff --git a/tests/trace_test_data/only_compressed_file1.dat.gz b/tests/trace_test_data/only_compressed_file1.dat.gz
new file mode 100644
index 0000000..4b8fa2c
--- /dev/null
+++ b/tests/trace_test_data/only_compressed_file1.dat.gz