diff options
| -rw-r--r-- | qolab/data/trace.py | 27 |
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 |