Merge pull request #9 from matthew-brett/small-streamlines-edits

RF: small streamlines API PR edits

Author: MarcCote

nibabel/streamlines/tests/test_trk.py

@@ -7,7 +7,7 @@
 from nibabel.externals.six import BytesIO
 
 from nibabel.testing import data_path
-from nibabel.testing import clear_and_catch_warnings
+from nibabel.testing import clear_and_catch_warnings, assert_arr_dict_equal
 from nose.tools import assert_equal, assert_raises, assert_true
 from numpy.testing import assert_array_equal
 
@@ -16,7 +16,7 @@
 from ..tractogram_file import HeaderError, HeaderWarning
 
 from .. import trk as trk_module
-from ..trk import TrkFile
+from ..trk import TrkFile, encode_value_in_name, decode_value_from_name
 from ..header import Field
 
 DATA = {}
@@ -82,14 +82,6 @@ def setup():
                                             affine_to_rasmm=np.eye(4))
 
 
-def assert_header_equal(h1, h2):
-    for k in h1.keys():
-        assert_array_equal(h2[k], h1[k])
-
-    for k in h2.keys():
-        assert_array_equal(h1[k], h2[k])
-
-
 class TestTRK(unittest.TestCase):
 
     def test_load_empty_file(self):
@@ -178,10 +170,7 @@ def test_load_complex_file_in_big_endian(self):
     def test_tractogram_file_properties(self):
         trk = TrkFile.load(DATA['simple_trk_fname'])
         assert_equal(trk.streamlines, trk.tractogram.streamlines)
-        assert_equal(trk.get_streamlines(), trk.streamlines)
-        assert_equal(trk.get_tractogram(), trk.tractogram)
-        assert_equal(trk.get_header(), trk.header)
-        assert_array_equal(trk.get_affine(), trk.header[Field.VOXEL_TO_RASMM])
+        assert_array_equal(trk.affine, trk.header[Field.VOXEL_TO_RASMM])
 
     def test_write_empty_file(self):
         tractogram = Tractogram(affine_to_rasmm=np.eye(4))
@@ -298,7 +287,7 @@ def test_load_write_LPS_file(self):
 
         new_trk = TrkFile.load(trk_file)
 
-        assert_header_equal(new_trk.header, trk.header)
+        assert_arr_dict_equal(new_trk.header, trk.header)
         assert_tractogram_equal(new_trk.tractogram, trk.tractogram)
 
         new_trk_orig = TrkFile.load(DATA['standard_LPS_trk_fname'])
@@ -322,7 +311,7 @@ def test_load_write_LPS_file(self):
 
         new_trk = TrkFile.load(trk_file)
 
-        assert_header_equal(new_trk.header, trk_LPS.header)
+        assert_arr_dict_equal(new_trk.header, trk_LPS.header)
         assert_tractogram_equal(new_trk.tractogram, trk.tractogram)
 
         new_trk_orig = TrkFile.load(DATA['standard_LPS_trk_fname'])
@@ -460,3 +449,47 @@ def test_write_scalars_and_properties_name_too_long(self):
     def test_str(self):
         trk = TrkFile.load(DATA['complex_trk_fname'])
         str(trk)  # Simply test it's not failing when called.
+
+    def test_header_read_restore(self):
+        # Test that reading a header restores the file position
+        trk_fname = DATA['simple_trk_fname']
+        bio = BytesIO()
+        bio.write(b'Along my very merry way')
+        hdr_pos = bio.tell()
+        hdr_from_fname = TrkFile._read_header(trk_fname)
+        with open(trk_fname, 'rb') as fobj:
+            bio.write(fobj.read())
+        bio.seek(hdr_pos)
+        # Check header is as expected
+        hdr_from_fname['_offset_data'] += hdr_pos  # Correct for start position
+        assert_arr_dict_equal(TrkFile._read_header(bio), hdr_from_fname)
+        # Check fileobject file position has not changed
+        assert_equal(bio.tell(), hdr_pos)
+
+
+def test_encode_names():
+    # Test function for encoding numbers into property names
+    b0 = b'\x00'
+    assert_equal(encode_value_in_name(0, 'foo', 10),
+                 b'foo' + b0 * 7)
+    assert_equal(encode_value_in_name(1, 'foo', 10),
+                 b'foo' + b0 * 7)
+    assert_equal(encode_value_in_name(8, 'foo', 10),
+                 b'foo' + b0 + b'8' + b0 * 5)
+    assert_equal(encode_value_in_name(40, 'foobar', 10),
+                 b'foobar' + b0 + b'40' + b0)
+    assert_equal(encode_value_in_name(1, 'foobarbazz', 10), b'foobarbazz')
+    assert_raises(ValueError, encode_value_in_name, 1, 'foobarbazzz', 10)
+    assert_raises(ValueError, encode_value_in_name, 2, 'foobarbaz', 10)
+    assert_equal(encode_value_in_name(2, 'foobarba', 10), b'foobarba\x002')
+
+
+def test_decode_names():
+    # Test function for decoding name string into name, number
+    b0 = b'\x00'
+    assert_equal(decode_value_from_name(b''), ('', 0))
+    assert_equal(decode_value_from_name(b'foo' + b0 * 7), ('foo', 1))
+    assert_equal(decode_value_from_name(b'foo\x008' + b0 * 5), ('foo', 8))
+    assert_equal(decode_value_from_name(b'foobar\x0010\x00'), ('foobar', 10))
+    assert_raises(ValueError, decode_value_from_name, b'foobar\x0010\x01')
+    assert_raises(HeaderError, decode_value_from_name, b'foo\x0010\x00111')

nibabel/streamlines/tractogram.py

@@ -80,7 +80,7 @@ class PerArrayDict(SliceableDataDict):
 
     In addition, it makes sure the amount of data contained in those ndarrays
     matches the number of streamlines given at the instantiation of this
-    dictionary.
+    instance.
     """
     def __init__(self, nb_elements, *args, **kwargs):
         self.nb_elements = nb_elements
@@ -114,7 +114,7 @@ class PerArraySequenceDict(SliceableDataDict):
 
     In addition, it makes sure the amount of data contained in those array
     sequences matches the number of elements given at the instantiation
-    of this dictionary.
+    of the instance.
     """
     def __init__(self, nb_elements, *args, **kwargs):
         self.nb_elements = nb_elements
@@ -136,9 +136,9 @@ def __setitem__(self, key, value):
 class LazyDict(collections.MutableMapping):
     """ Dictionary of generator functions.
 
-    This container behaves like an dictionary but it makes sure its elements
-    are callable objects and assumed to be generator function yielding values.
-    When getting the element associated to a given key, the element (i.e. a
+    This container behaves like a dictionary but it makes sure its elements are
+    callable objects that it assumes are generator functions yielding values.
+    When getting the element associated with a given key, the element (i.e. a
     generator function) is first called before being returned.
     """
     def __init__(self, *args, **kwargs):
@@ -178,7 +178,7 @@ def __len__(self):
 class TractogramItem(object):
     """ Class containing information about one streamline.
 
-    :class:`TractogramItem` objects have three main properties: `streamline`,
+    :class:`TractogramItem` objects have three public attributes: `streamline`,
     `data_for_streamline`, and `data_for_points`.
 
     Parameters
@@ -187,14 +187,14 @@ class TractogramItem(object):
         Points of this streamline represented as an ndarray of shape (N, 3)
         where N is the number of points.
     data_for_streamline : dict
-        Dictionary containing some data associated to this particular
-        streamline. Each key `k` is mapped to a ndarray of shape (Pt,), where
-        `Pt` is the dimension of the data associated with key `k`.
+        Dictionary containing some data associated with this particular
+        streamline. Each key ``k`` is mapped to a ndarray of shape (Pt,), where
+        ``Pt`` is the dimension of the data associated with key ``k``.
     data_for_points : dict
         Dictionary containing some data associated to each point of this
-        particular streamline. Each key `k` is mapped to a ndarray of
-        shape (Nt, Mk), where `Nt` is the number of points of this streamline
-        and `Mk` is the dimension of the data associated with key `k`.
+        particular streamline. Each key ``k`` is mapped to a ndarray of shape
+        (Nt, Mk), where ``Nt`` is the number of points of this streamline and
+        ``Mk`` is the dimension of the data associated with key ``k``.
     """
     def __init__(self, streamline, data_for_streamline, data_for_points):
         self.streamline = np.asarray(streamline)
@@ -215,7 +215,7 @@ class Tractogram(object):
     choice as long as you provide the correct `affine_to_rasmm` matrix, at
     construction time, that brings the streamlines back to *RAS+*, *mm* space,
     where the coordinates (0,0,0) corresponds to the center of the voxel
-    (opposed to a corner).
+    (as opposed to the corner of the voxel).
 
     Attributes
     ----------
@@ -224,18 +224,18 @@ class Tractogram(object):
         shape ($N_t$, 3) where $N_t$ is the number of points of
         streamline $t$.
     data_per_streamline : :class:`PerArrayDict` object
-        Dictionary where the items are (str, 2D array).
-        Each key represents an information $i$ to be kept alongside every
-        streamline, and its associated value is a 2D array of shape
-        ($T$, $P_i$) where $T$ is the number of streamlines and $P_i$ is
-        the number of values to store for that particular information $i$.
+        Dictionary where the items are (str, 2D array).  Each key represents a
+        piece of information $i$ to be kept alongside every streamline, and its
+        associated value is a 2D array of shape ($T$, $P_i$) where $T$ is the
+        number of streamlines and $P_i$ is the number of values to store for
+        that particular piece of information $i$.
     data_per_point : :class:`PerArraySequenceDict` object
-        Dictionary where the items are (str, :class:`ArraySequence`).
-        Each key represents an information $i$ to be kept alongside every
-        point of every streamline, and its associated value is an iterable
-        of ndarrays of shape ($N_t$, $M_i$) where $N_t$ is the number of
-        points for a particular streamline $t$ and $M_i$ is the number
-        values to store for that particular information $i$.
+        Dictionary where the items are (str, :class:`ArraySequence`).  Each key
+        represents a piece of information $i$ to be kept alongside every point
+        of every streamline, and its associated value is an iterable of
+        ndarrays of shape ($N_t$, $M_i$) where $N_t$ is the number of points
+        for a particular streamline $t$ and $M_i$ is the number values to store
+        for that particular piece of information $i$.
     """
     def __init__(self, streamlines=None,
                  data_per_streamline=None,
@@ -424,29 +424,29 @@ class LazyTractogram(Tractogram):
     choice as long as you provide the correct `affine_to_rasmm` matrix, at
     construction time, that brings the streamlines back to *RAS+*, *mm* space,
     where the coordinates (0,0,0) corresponds to the center of the voxel
-    (opposed to a corner).
+    (as opposed to the corner of the voxel).
 
     Attributes
     ----------
     streamlines : generator function
         Generator function yielding streamlines. Each streamline is an
         ndarray of shape ($N_t$, 3) where $N_t$ is the number of points of
         streamline $t$.
-    data_per_streamline : :class:`LazyDict` object
+    data_per_streamline : instance of :class:`LazyDict`
         Dictionary where the items are (str, instantiated generator).
-        Each key represents an information $i$ to be kept alongside every
-        streamline, and its associated value is a generator function
-        yielding that information via ndarrays of shape ($P_i$,) where
-        $P_i$ is the number of values to store for that particular
-        information $i$.
+        Each key represents a piece of information $i$ to be kept alongside
+        every streamline, and its associated value is a generator function
+        yielding that information via ndarrays of shape ($P_i$,) where $P_i$ is
+        the number of values to store for that particular piece of information
+        $i$.
     data_per_point : :class:`LazyDict` object
-        Dictionary where the items are (str, instantiated generator).
-        Each key represents an information $i$ to be kept alongside every
-        point of every streamline, and its associated value is a generator
-        function yielding that information via ndarrays of shape
-        ($N_t$, $M_i$) where $N_t$ is the number of points for a particular
-        streamline $t$ and $M_i$ is the number of values to store for
-        that particular information $i$.
+        Dictionary where the items are (str, instantiated generator).  Each key
+        represents a piece of information $i$ to be kept alongside every point
+        of every streamline, and its associated value is a generator function
+        yielding that information via ndarrays of shape ($N_t$, $M_i$) where
+        $N_t$ is the number of points for a particular streamline $t$ and $M_i$
+        is the number of values to store for that particular piece of
+        information $i$.
 
     Notes
     -----
@@ -599,7 +599,6 @@ def _apply_affine():
     def _set_streamlines(self, value):
         if value is not None and not callable(value):
             raise TypeError("`streamlines` must be a generator function.")
-
         self._streamlines = value
 
     @property

nibabel/streamlines/tractogram_file.py

@@ -1,3 +1,5 @@
+"""  Define abstract interface for Tractogram file classes
+"""
 from abc import ABCMeta, abstractmethod
 from nibabel.externals.six import with_metaclass
 
@@ -49,21 +51,9 @@ def header(self):
 
     @property
     def affine(self):
+        """ voxmm -> rasmm affine. """
         return self.header.get(Field.VOXEL_TO_RASMM)
 
-    def get_tractogram(self):
-        return self.tractogram
-
-    def get_streamlines(self):
-        return self.streamlines
-
-    def get_header(self):
-        return self.header
-
-    def get_affine(self):
-        """ Returns vox -> rasmm affine. """
-        return self.affine
-
     @abstractclassmethod
     def is_correct_format(cls, fileobj):
         """ Checks if the file has the right streamlines file format.
@@ -85,7 +75,7 @@ def is_correct_format(cls, fileobj):
 
     @abstractclassmethod
     def load(cls, fileobj, lazy_load=True):
-        """ Loads streamlines from a file-like object.
+        """ Loads streamlines from a filename or file-like object.
 
         Parameters
         ----------
@@ -107,7 +97,7 @@ def load(cls, fileobj, lazy_load=True):
 
     @abstractmethod
     def save(self, fileobj):
-        """ Saves streamlines to a file-like object.
+        """ Saves streamlines to a filename or file-like object.
 
         Parameters
         ----------