diff --git a/nibabel/streamlines/tck.py b/nibabel/streamlines/tck.py index ec8e7dbce7..43df2f87e0 100644 --- a/nibabel/streamlines/tck.py +++ b/nibabel/streamlines/tck.py @@ -6,6 +6,7 @@ import os import warnings +from contextlib import suppress import numpy as np @@ -266,11 +267,6 @@ def _write_header(fileobj, header): ) out = '\n'.join(lines) - # Check the header is well formatted. - if out.count('\n') > len(lines) - 1: # \n only allowed between lines. - msg = f"Key-value pairs cannot contain '\\n':\n{out}" - raise HeaderError(msg) - if out.count(':') > len(lines): # : only one per line (except the last one which contains END). msg = f"Key-value pairs cannot contain ':':\n{out}" @@ -331,6 +327,8 @@ def _read_header(cls, fileobj): f.seek(1, os.SEEK_CUR) # Skip \n found_end = False + key = None + tmp_hdr = {} # Read all key-value pairs contained in the header, stop at EOF for n_line, line in enumerate(f, 1): @@ -343,15 +341,22 @@ def _read_header(cls, fileobj): found_end = True break - if ':' not in line: # Invalid header line + # Set new key if available, otherwise append to last known key + with suppress(ValueError): + key, line = line.split(':', 1) + key = key.strip() + + # Apparent continuation line before any keys are found + if key is None: raise HeaderError(f'Invalid header (line {n_line}): {line}') - key, value = line.split(':', 1) - hdr[key.strip()] = value.strip() + tmp_hdr.setdefault(key, []).append(line.strip()) if not found_end: raise HeaderError('Missing END in the header.') + hdr.update({key: '\n'.join(val) for key, val in tmp_hdr.items()}) + offset_data = f.tell() # Set the file position where it was, in case it was previously open diff --git a/nibabel/streamlines/tests/test_tck.py b/nibabel/streamlines/tests/test_tck.py index f514d3f3df..3df7dd4f2d 100644 --- a/nibabel/streamlines/tests/test_tck.py +++ b/nibabel/streamlines/tests/test_tck.py @@ -31,6 +31,7 @@ def setup_module(): # standard.tck contains only streamlines DATA['standard_tck_fname'] = pjoin(data_path, 'standard.tck') DATA['matlab_nan_tck_fname'] = pjoin(data_path, 'matlab_nan.tck') + DATA['multiline_header_fname'] = pjoin(data_path, 'multiline_header_field.tck') DATA['streamlines'] = [ np.arange(1 * 3, dtype='f4').reshape((1, 3)), @@ -87,6 +88,14 @@ def test_load_matlab_nan_file(self): assert len(streamlines) == 1 assert streamlines[0].shape == (108, 3) + def test_load_multiline_header_file(self): + for lazy_load in [False, True]: + tck = TckFile.load(DATA['multiline_header_fname'], lazy_load=lazy_load) + streamlines = list(tck.tractogram.streamlines) + assert len(tck.header['command_history'].splitlines()) == 3 + assert len(streamlines) == 1 + assert streamlines[0].shape == (253, 3) + def test_writeable_data(self): data = DATA['simple_tractogram'] for key in ('simple_tck_fname', 'simple_tck_big_endian_fname'): @@ -192,10 +201,6 @@ def test_write_simple_file(self): # TCK file containing not well formatted entries in its header. tck_file = BytesIO() tck = TckFile(tractogram) - tck.header['new_entry'] = 'value\n' # \n not allowed - with pytest.raises(HeaderError): - tck.save(tck_file) - tck.header['new_entry'] = 'val:ue' # : not allowed with pytest.raises(HeaderError): tck.save(tck_file) diff --git a/nibabel/tests/data/multiline_header_field.tck b/nibabel/tests/data/multiline_header_field.tck new file mode 100644 index 0000000000..42ebedc43a Binary files /dev/null and b/nibabel/tests/data/multiline_header_field.tck differ