BF: fixed TCK buferring and added a unit test

Author: MarcCote

nibabel/streamlines/tck.py

@@ -347,9 +347,11 @@ def _read(fileobj, header, buffer_size=4):
         points : ndarray of shape (n_pts, 3)
             Streamline points
         """
-        buffer_size = buffer_size * MEGABYTE
-        buffer_size += 3 - (buffer_size % 3)  # Make it a multiple of 3.
         dtype = header["_dtype"]
+        coordinate_size = 3 * dtype.itemsize
+        # Make buffer_size an integer and a multiple of coordinate_size.
+        buffer_size = int(buffer_size * MEGABYTE)
+        buffer_size += coordinate_size - (buffer_size % coordinate_size)
 
         with Opener(fileobj) as f:
             start_position = f.tell()

nibabel/streamlines/tests/test_tck.py

@@ -10,6 +10,7 @@
 
 from nibabel.testing import data_path
 from .test_tractogram import assert_tractogram_equal
+from ..array_sequence import ArraySequence
 from ..tractogram import Tractogram
 from ..tractogram_file import DataError
 
@@ -80,13 +81,23 @@ class TestTCK(unittest.TestCase):
 
     def test_load_empty_file(self):
         for lazy_load in [False, True]:
-            trk = TckFile.load(DATA['empty_tck_fname'], lazy_load=lazy_load)
-            assert_tractogram_equal(trk.tractogram, DATA['empty_tractogram'])
+            tck = TckFile.load(DATA['empty_tck_fname'], lazy_load=lazy_load)
+            assert_tractogram_equal(tck.tractogram, DATA['empty_tractogram'])
 
     def test_load_simple_file(self):
         for lazy_load in [False, True]:
-            trk = TckFile.load(DATA['simple_tck_fname'], lazy_load=lazy_load)
-            assert_tractogram_equal(trk.tractogram, DATA['simple_tractogram'])
+            tck = TckFile.load(DATA['simple_tck_fname'], lazy_load=lazy_load)
+            assert_tractogram_equal(tck.tractogram, DATA['simple_tractogram'])
+
+        # Force TCK loading to use buffering.
+        buffer_size = 1. / 1024**2  # 1 bytes
+        hdr = TckFile._read_header(DATA['simple_tck_fname'])
+        tck_reader = TckFile._read(DATA['simple_tck_fname'], hdr, buffer_size)
+        streamlines = ArraySequence(tck_reader)
+        tractogram = Tractogram(streamlines)
+        tractogram.affine_to_rasmm = np.eye(4)
+        tck = TckFile(tractogram, header=hdr)
+        assert_tractogram_equal(tck.tractogram, DATA['simple_tractogram'])
 
     def test_load_simple_file_in_big_endian(self):
         for lazy_load in [False, True]: