diff --git a/corsikaio/file.py b/corsikaio/file.py index b42f04b..c31b5d7 100644 --- a/corsikaio/file.py +++ b/corsikaio/file.py @@ -22,7 +22,6 @@ ParticleEvent = namedtuple('ParticleEvent', ['header', 'particles', 'longitudinal', 'end']) - class CorsikaFile: def __init__(self, path): @@ -67,18 +66,15 @@ def __next__(self): self._run_end = parse_run_end(block) raise StopIteration() - if len(block) < BLOCK_SIZE_BYTES: - raise StopIteration - if block[:4] != b'EVTH': raise IOError('EVTH block expected but found {}'.format(block[:4])) event_header = parse_event_header(block)[0] - block = self.read_block() data_bytes = bytearray() long_bytes = bytearray() + block = self.read_block() while block[:4] != b'EVTE': if block[:4] == b'LONG': diff --git a/corsikaio/io.py b/corsikaio/io.py index b7a9231..78be7b3 100644 --- a/corsikaio/io.py +++ b/corsikaio/io.py @@ -70,4 +70,8 @@ def read_block(f, buffer_size=None): if (pos + 4) % (buffer_size + 8) == 0: f.read(8) - return f.read(BLOCK_SIZE_BYTES) + block = f.read(BLOCK_SIZE_BYTES) + if len(block) < BLOCK_SIZE_BYTES: + raise IOError("Read less bytes than expected, file seems to be truncated") + + return block diff --git a/tests/test_file.py b/tests/test_file.py index e5424bb..cdfddb1 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -1,3 +1,6 @@ +import pytest + + def test_version(): from corsikaio import CorsikaFile @@ -66,3 +69,23 @@ def test_particle_longi(): e = next(f) assert e.header['event_number'] == 1 + + +def test_truncated(tmp_path): + '''Test we raise a meaningful error for a truncated file + + Truncated files might happen if corsika crashes or the disk is full. + + Regression test for cta-observatory/pycorsikaio#15 + ''' + from corsikaio import CorsikaParticleFile + path = tmp_path / "truncated.dat" + + with open("tests/resources/corsika757_particle", "rb") as f: + with path.open("wb") as out: + out.write(f.read(273 * 10)) + + with pytest.raises(IOError, match="seems to be truncated"): + with CorsikaParticleFile(path) as f: + for _ in f: + pass