diff options
| -rw-r--r-- | qolab/data/trace.py | 37 | ||||
| -rw-r--r-- | qolab/file_utils/__init__.py | 10 |
2 files changed, 26 insertions, 21 deletions
diff --git a/qolab/data/trace.py b/qolab/data/trace.py index bda4da5..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 @@ -25,28 +25,23 @@ def loadTraceRawHeaderAndData(fname, tryCompressedIfMissing=True): headerstr=[] data = None 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' + # attempt to locate compressed file for the missing base file + for ext in ['gz', 'bz', 'bz2']: + if os.path.exists(fname+'.'+ext): + fname += '.'+ext + 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 diff --git a/qolab/file_utils/__init__.py b/qolab/file_utils/__init__.py index e61811e..bcc0824 100644 --- a/qolab/file_utils/__init__.py +++ b/qolab/file_utils/__init__.py @@ -69,6 +69,16 @@ def get_next_data_file(prefix, savepath, run_number=None, datestr=None, date_for fname = os.path.join(savepath, f'{prefix}_{datestr}_{run_number:05d}.{extension}') return(fname) +def infer_compression(fname): + """Infers compression algorithm from filename extension""" + compression = None # usual suspect + b, fext = os.path.splitext(fname) + if fext == '.gz': + compression = 'gzip' + elif ( fext == '.bz') or (fext == '.bz2'): + compression = 'bzip' + return compression + def save_table_with_header(fname, data, header='', comment_symbol='%', skip_headers_if_file_exist=False, item_format='e', item_separator='\t', compressionmethod=None, compresslevel=9): """Saves output to CSV or TSV file with specially formatted header. |