Addressed @matthew-brett's comments

Author: MarcCote

nibabel/streamlines/tests/test_tractogram.py

@@ -3,7 +3,7 @@
 import numpy as np
 import warnings
 
-from nibabel.testing import assert_arrays_equal, check_iteration
+from nibabel.testing import assert_arrays_equal
 from nibabel.testing import clear_and_catch_warnings
 from nose.tools import assert_equal, assert_raises, assert_true
 from numpy.testing import assert_array_equal, assert_array_almost_equal
@@ -110,7 +110,7 @@ def check_tractogram(tractogram,
     streamlines = list(streamlines)
     assert_equal(len(tractogram), len(streamlines))
     assert_arrays_equal(tractogram.streamlines, streamlines)
-    assert_true(check_iteration(tractogram))
+    [t for t in tractogram]  # Force iteration through tractogram.
 
     assert_equal(len(tractogram.data_per_streamline), len(data_per_streamline))
     for key in data_per_streamline.keys():
@@ -556,7 +556,7 @@ def test_lazy_tractogram_creation(self):
                                     DATA['data_per_streamline_func'],
                                     DATA['data_per_point_func'])
 
-        assert_true(check_iteration(tractogram))
+        [t for t in tractogram]  # Force iteration through tractogram.
         assert_equal(len(tractogram), len(DATA['streamlines']))
 
         # Generator functions get re-called and creates new iterators.
@@ -627,7 +627,7 @@ def test_lazy_tractogram_len(self):
             tractogram = LazyTractogram(DATA['streamlines_func'])
 
             assert_true(tractogram._nb_streamlines is None)
-            check_iteration(tractogram)  # Force iteration through tractogram.
+            [t for t in tractogram]  # Force iteration through tractogram.
             assert_equal(tractogram._nb_streamlines, len(DATA['streamlines']))
             # This should *not* produce a warning.
             assert_equal(len(tractogram), len(DATA['streamlines']))

nibabel/streamlines/trk.py

@@ -100,8 +100,8 @@ def get_affine_trackvis_to_rasmm(header):
     affine = np.dot(scale, affine)
 
     # TrackVis considers coordinate (0,0,0) to be the corner of the
-    # voxel whereas streamlines returned assume (0,0,0) to be the
-    # center of the voxel. Thus, streamlines are shifted of half a voxel.
+    # voxel whereas streamlines returned assumes (0,0,0) to be the
+    # center of the voxel. Thus, streamlines are shifted by half a voxel.
     offset = np.eye(4)
     offset[:-1, -1] -= 0.5
     affine = np.dot(offset, affine)
@@ -132,21 +132,24 @@ def encode_value_in_name(value, name, max_name_len=20):
     """ Encodes a value in the last two bytes of a string.
 
     If `value` is one, then there is no encoding and the last two bytes
-    are left untouched. This function also verify that the length of name is
-    less than `max_name_len`.
+    are left untouched. Otherwise, the byte before the last will be
+    set to \x00 and the last byte will correspond to the value.
+
+    This function also verifies that the length of name is less
+    than `max_name_len`.
 
     Parameters
     ----------
-    value : int
-        Integer value to encode.
-    name : str
+    value : byte
+        Integer value between 0 and 255 to encode.
+    name : bytes
         Name in which the last two bytes will serve to encode `value`.
     max_name_len : int, optional
         Maximum length name can have.
 
     Returns
     -------
-    encoded_name : str
+    encoded_name : bytes
         Name containing the encoded value.
     """
 
@@ -161,9 +164,10 @@ def encode_value_in_name(value, name, max_name_len=20):
                ).format(name, max_name_len - 2)
         raise ValueError(msg)
 
+    name = name.ljust(max_name_len, '\x00')
     if value > 1:
         # Use the last two bytes of `name` to store `value`.
-        name = (asbytes(name[:18].ljust(18, '\x00')) + b'\x00' +
+        name = (asbytes(name[:max_name_len - 2]) + b'\x00' +
                 np.array(value, dtype=np.int8).tostring())
 
     return name
@@ -183,10 +187,10 @@ class TrkReader(object):
     ----
     TrackVis (so its file format: TRK) considers the streamline coordinate
     (0,0,0) to be in the corner of the voxel whereas NiBabel's streamlines
-    internal representation (Voxel space) assume (0,0,0) to be in the
+    internal representation (Voxel space) assumes (0,0,0) to be in the
     center of the voxel.
 
-    Thus, streamlines are shifted of half a voxel on load and are shifted
+    Thus, streamlines are shifted by half a voxel on load and are shifted
     back on save.
     """
     def __init__(self, fileobj):
@@ -461,10 +465,10 @@ class TrkFile(TractogramFile):
     ----
     TrackVis (so its file format: TRK) considers the streamline coordinate
     (0,0,0) to be in the corner of the voxel whereas NiBabel's streamlines
-    internal representation (Voxel space) assume (0,0,0) to be in the
+    internal representation (Voxel space) assumes (0,0,0) to be in the
     center of the voxel.
 
-    Thus, streamlines are shifted of half a voxel on load and are shifted
+    Thus, streamlines are shifted by half a voxel on load and are shifted
     back on save.
     """
 
@@ -515,8 +519,9 @@ def is_correct_format(cls, fileobj):
             otherwise returns False.
         """
         with Opener(fileobj) as f:
-            magic_number = f.read(5)
-            f.seek(-5, os.SEEK_CUR)
+            magic_len = len(cls.MAGIC_NUMBER)
+            magic_number = f.read(magic_len)
+            f.seek(-magic_len, os.SEEK_CUR)
             return magic_number == cls.MAGIC_NUMBER
 
     @classmethod

nibabel/streamlines/utils.py

@@ -1,4 +1,3 @@
-import numpy as np
 import nibabel
 
 
@@ -9,7 +8,7 @@ def get_affine_from_reference(ref):
     ---------
     ref : str or :class:`Nifti1Image` object or ndarray shape (4, 4)
         If str then it's the filename of reference file that will be loaded
-        using :func:nibabel.load in order to obtain the affine.
+        using :func:`nibabel.load` in order to obtain the affine.
         If :class:`Nifti1Image` object then the affine is obtained from it.
         If ndarray shape (4, 4) then it's the affine.
 
@@ -18,14 +17,15 @@ def get_affine_from_reference(ref):
     affine : ndarray (4, 4)
         Transformation matrix mapping voxel space to RAS+mm space.
     """
-    if type(ref) is np.ndarray:
+    if hasattr(ref, 'affine'):
+        return ref.affine
+
+    if hasattr(ref, 'shape'):
         if ref.shape != (4, 4):
             msg = "`ref` needs to be a numpy array with shape (4, 4)!"
             raise ValueError(msg)
 
         return ref
-    elif hasattr(ref, 'affine'):
-        return ref.affine
 
     # Assume `ref` is the name of a neuroimaging file.
     return nibabel.load(ref).affine

nibabel/testing/__init__.py

@@ -65,19 +65,8 @@ def assert_allclose_safely(a, b, match_nans=True, rtol=1e-5, atol=1e-8):
     assert_true(np.allclose(a, b, rtol=rtol, atol=atol))
 
 
-def check_iteration(iterable):
-    """ Checks that an object can be iterated through without errors. """
-    try:
-        for _ in iterable:
-            pass
-    except:
-        return False
-
-    return True
-
-
 def assert_arrays_equal(arrays1, arrays2):
-    """ Checks that two iterables yielding arrays are equals. """
+    """ Check two iterables yield the same sequence of arrays. """
     for arr1, arr2 in zip_longest(arrays1, arrays2, fillvalue=None):
         assert_false(arr1 is None or arr2 is None)
         assert_array_equal(arr1, arr2)