TST: added unit test for TCK files written in big-endian

Author: MarcCote

nibabel/streamlines/tck.py

@@ -14,7 +14,7 @@
 
 from .array_sequence import ArraySequence
 from .tractogram_file import TractogramFile
-from .tractogram_file import HeaderError, DataWarning
+from .tractogram_file import HeaderError, DataWarning, DataError
 from .tractogram import TractogramItem, Tractogram, LazyTractogram
 from .header import Field
 
@@ -206,8 +206,8 @@ def save(self, fileobj):
 
             for t in tractogram:
                 s = t.streamline.astype(dtype)
-                f.write(asbytes(np.r_[s.astype('<f4'),
-                                      self.FIBER_DELIMITER].tostring()))
+                data = np.r_[s, self.FIBER_DELIMITER]
+                f.write(asbytes(data.tostring()))
                 nb_streamlines += 1
 
             header[Field.NB_STREAMLINES] = nb_streamlines
@@ -281,7 +281,7 @@ def _read_header(fileobj):
 
         with Opener(fileobj) as f:
             # Read magic number
-            buf = asstr(f.fobj.readline())
+            magic_number = asstr(f.fobj.readline())
 
             # Read all key-value pairs contained in the header.
             buf = asstr(f.fobj.readline())
@@ -290,6 +290,7 @@ def _read_header(fileobj):
 
         # Build header dictionary from the buffer.
         hdr = dict(item.split(': ') for item in buf.rstrip().split('\n')[:-1])
+        hdr[Field.MAGIC_NUMBER] = magic_number
 
         # Check integrity of TCK header.
         if 'datatype' not in hdr:
@@ -379,7 +380,15 @@ def _read(fileobj, header, buffer_size=4):
                 pts = pts.reshape([-1, 3])
                 idx_nan = np.arange(len(pts))[np.isnan(pts[:, 0])]
 
+                # Make sure we've read enough to find a streamline delimiter.
                 if len(idx_nan) == 0:
+                    # If we've read the whole file, then fail.
+                    if eof and not np.all(np.isinf(pts)):
+                        msg = ("Cannot find a streamline delimiter. This file"
+                               " might be corrupted.")
+                        raise DataError(msg)
+
+                    # Otherwise read a bit more.
                     continue
 
                 nb_pts_total = 0

nibabel/streamlines/tests/test_tck.py

@@ -5,11 +5,12 @@
 
 from nibabel.externals.six import BytesIO
 
-from nose.tools import assert_equal
+from nose.tools import assert_equal, assert_raises
 
 from nibabel.testing import data_path
 from .test_tractogram import assert_tractogram_equal
 from ..tractogram import Tractogram
+from ..tractogram_file import DataError
 
 from ..tck import TckFile
 
@@ -86,6 +87,20 @@ def test_load_simple_file(self):
             trk = TckFile.load(DATA['simple_tck_fname'], lazy_load=lazy_load)
             assert_tractogram_equal(trk.tractogram, DATA['simple_tractogram'])
 
+    def test_load_simple_file_in_big_endian(self):
+        for lazy_load in [False, True]:
+            tck = TckFile.load(DATA['simple_tck_big_endian_fname'],
+                               lazy_load=lazy_load)
+            assert_tractogram_equal(tck.tractogram, DATA['simple_tractogram'])
+            assert_equal(tck.header['datatype'], 'Float32BE')
+
+    def test_load_file_with_wrong_information(self):
+        tck_file = open(DATA['simple_tck_fname'], 'rb').read()
+
+        # Simulate a TCK file where `datatype` was incorrectly provided.
+        new_tck_file = tck_file.replace("Float32LE", "Float32BE")
+        assert_raises(DataError, TckFile.load, BytesIO(new_tck_file))
+
     def test_write_empty_file(self):
         tractogram = Tractogram(affine_to_rasmm=np.eye(4))
 
@@ -149,3 +164,7 @@ def test_load_write_file(self):
             loaded_tck = TckFile.load(tck_file, lazy_load=False)
             assert_tractogram_equal(loaded_tck.tractogram,
                                     tractogram.to_world(lazy=True))
+
+    def test_str(self):
+        tck = TckFile.load(DATA['simple_tck_fname'])
+        str(tck)  # Simply test it's not failing when called.