diff --git a/Semi_ATE/STDF/utils.py b/Semi_ATE/STDF/utils.py index fcee5da..e5e722b 100644 --- a/Semi_ATE/STDF/utils.py +++ b/Semi_ATE/STDF/utils.py @@ -721,26 +721,34 @@ def __init__(self, FileName, unpack=False, of_interest=None): raise STDFError("'%s' does not exist" %(FileName)) # seimit : adding support for compressed files compression = extension_from_magic_number_in_file(FileName) - if compression[0] == '.xz': + # `extension_from_magic_number_in_file` returns an empty list + # when no magic signature is recognised (e.g. some bare-binary + # STDF files produced by older ATE tooling). Indexing + # `compression[0]` unconditionally raises IndexError in that + # case and prevents the file from ever being opened — see #75. + compression_ext = compression[0] if compression else None + if compression_ext == '.xz': import lzma self.fd = lzma.open(FileName, 'rb') self.parse_FAR() - elif compression[0] == '.bz2': + elif compression_ext == '.bz2': import bz2 self.fd = bz2.open(FileName, 'rb') self.parse_FAR() - elif compression[0] == '.gz': + elif compression_ext == '.gz': import gzip self.fd = gzip.open(FileName, 'rb') self.parse_FAR() - elif compression[0] == '.zip': + elif compression_ext == '.zip': import zipfile zfile = zipfile.ZipFile(FileName, 'r') for name in zfile.namelist(): self.fd = zfile.open(name) self.parse_FAR() else: - # Assume standard binary stdf file + # Assume standard binary stdf file (no recognised + # compression magic, or magic that does not match any + # supported compression scheme). self.endian = get_STDF_setup_from_file(FileName)[0] self.version = 'V%s' % struct.unpack( 'B', get_bytes_from_file(FileName, 5, 1)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 82d1095..3334c66 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -203,4 +203,23 @@ def test_dict_to_rec(): assert rec.get_value('HEAD_NUM') == head_num assert rec.get_value('SITE_GRP') == site_grp assert rec.get_value('START_T') == start_t - assert rec.get_value('WAFER_ID') == waf_id \ No newline at end of file + assert rec.get_value('WAFER_ID') == waf_id + +def test_records_from_file_handles_unrecognised_magic_number(): + """Regression for #75: a bare-binary STDF file whose magic number does + not match any recognised signature must not crash records_from_file + with IndexError. The previous code did ``compression[0]`` without + checking that the list was non-empty.""" + from unittest.mock import patch + + far_bytes = STDF.FAR().__repr__() + with tempfile.NamedTemporaryFile(mode="wb", suffix=".stdf", delete=False) as f: + f.write(far_bytes) + file_path = f.name + + with patch("Semi_ATE.STDF.utils.extension_from_magic_number_in_file", + return_value=[]): + records = list(STDF.records_from_file(file_path)) + + assert len(records) >= 1 + assert records[0].id == "FAR"