Merge pull request #1175 from effigies/fix-read-tckedit

ENH: Support multiline header fields in TCK

Author: effigies

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

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)