Addressed @matthew-brett's about coroutine and ArraySequence.extend method.

Author: MarcCote

nibabel/streamlines/__init__.py

@@ -127,6 +127,6 @@ def save(tractogram, filename, **kwargs):
         if len(kwargs) > 0:
             msg = ("A 'TractogramFile' object was provided, no need for"
                    " keyword arguments.")
-        raise ValueError(msg)
+            raise ValueError(msg)
 
     tractogram_file.save(filename)

nibabel/streamlines/array_sequence.py

@@ -1,3 +1,4 @@
+import numbers
 import numpy as np
 
 
@@ -37,7 +38,7 @@ def __init__(self, iterable=None):
         """
         # Create new empty `ArraySequence` object.
         self._is_view = False
-        self._data = np.array(0)
+        self._data = np.array([])
         self._offsets = np.array([], dtype=np.intp)
         self._lengths = np.array([], dtype=np.intp)
 
@@ -79,8 +80,7 @@ def __init__(self, iterable=None):
         self._lengths = np.asarray(lengths)
 
         # Clear unused memory.
-        if self._data.ndim != 0:
-            self._data.resize((offset,) + self.common_shape)
+        self._data.resize((offset,) + self.common_shape)
 
     @property
     def is_array_sequence(self):
@@ -89,13 +89,10 @@ def is_array_sequence(self):
     @property
     def common_shape(self):
         """ Matching shape of the elements in this array sequence. """
-        if self._data.ndim == 0:
-            return ()
-
         return self._data.shape[1:]
 
     def append(self, element):
-        """ Appends :obj:`element` to this array sequence.
+        """ Appends `element` to this array sequence.
 
         Parameters
         ----------
@@ -108,28 +105,28 @@ def append(self, element):
         If you need to add multiple elements you should consider
         `ArraySequence.extend`.
         """
-        if self._data.ndim == 0:
-            self._data = np.asarray(element).copy()
-            self._offsets = np.array([0])
-            self._lengths = np.array([len(element)])
-            return
+        element = np.asarray(element)
 
-        if element.shape[1:] != self.common_shape:
+        if self.common_shape != () and element.shape[1:] != self.common_shape:
             msg = "All dimensions, except the first one, must match exactly"
             raise ValueError(msg)
 
-        self._offsets = np.r_[self._offsets, len(self._data)]
-        self._lengths = np.r_[self._lengths, len(element)]
-        self._data = np.append(self._data, element, axis=0)
+        next_offset = self._data.shape[0]
+        size = (self._data.shape[0] + element.shape[0],) + element.shape[1:]
+        self._data.resize(size)
+        self._data[next_offset:] = element
+        self._offsets = np.r_[self._offsets, next_offset]
+        self._lengths = np.r_[self._lengths, element.shape[0]]
 
     def extend(self, elements):
         """ Appends all `elements` to this array sequence.
 
         Parameters
         ----------
-        elements : list of ndarrays or :class:`ArraySequence` object
-            If list of ndarrays, each ndarray will be concatenated along the
-            first dimension then appended to the data of this ArraySequence.
+        elements : iterable of ndarrays or :class:`ArraySequence` object
+            If iterable of ndarrays, each ndarray will be concatenated along
+            the first dimension then appended to the data of this
+            ArraySequence.
             If :class:`ArraySequence` object, its data are simply appended to
             the data of this ArraySequence.
 
@@ -138,35 +135,31 @@ def extend(self, elements):
             The shape of the elements to be added must match the one of the
             data of this :class:`ArraySequence` except for the first dimension.
         """
+        if not is_array_sequence(elements):
+            self.extend(ArraySequence(elements))
+            return
+
         if len(elements) == 0:
             return
 
-        if self._data.ndim == 0:
-            elem = np.asarray(elements[0])
-            self._data = np.zeros((0, elem.shape[1]), dtype=elem.dtype)
+        if (self.common_shape != () and
+                elements.common_shape != self.common_shape):
+            msg = "All dimensions, except the first one, must match exactly"
+            raise ValueError(msg)
 
         next_offset = self._data.shape[0]
+        self._data.resize((self._data.shape[0] + sum(elements._lengths),
+                           elements._data.shape[1]))
 
-        if is_array_sequence(elements):
-            self._data.resize((self._data.shape[0] + sum(elements._lengths),
-                               self._data.shape[1]))
-
-            offsets = []
-            for offset, length in zip(elements._offsets, elements._lengths):
-                offsets.append(next_offset)
-                chunk = elements._data[offset:offset + length]
-                self._data[next_offset:next_offset + length] = chunk
-                next_offset += length
-
-            self._lengths = np.r_[self._lengths, elements._lengths]
-            self._offsets = np.r_[self._offsets, offsets]
+        offsets = []
+        for offset, length in zip(elements._offsets, elements._lengths):
+            offsets.append(next_offset)
+            chunk = elements._data[offset:offset + length]
+            self._data[next_offset:next_offset + length] = chunk
+            next_offset += length
 
-        else:
-            self._data = np.concatenate([self._data] + list(elements), axis=0)
-            lengths = list(map(len, elements))
-            self._lengths = np.r_[self._lengths, lengths]
-            self._offsets = np.r_[self._offsets,
-                                  np.cumsum([next_offset] + lengths)[:-1]]
+        self._lengths = np.r_[self._lengths, elements._lengths]
+        self._offsets = np.r_[self._offsets, offsets]
 
     def copy(self):
         """ Creates a copy of this :class:`ArraySequence` object. """
@@ -210,7 +203,7 @@ def __getitem__(self, idx):
             Otherwise, returns a :class:`ArraySequence` object which is view
             of the selected sequences.
         """
-        if isinstance(idx, (int, np.integer)):
+        if isinstance(idx, (numbers.Integral, np.integer)):
             start = self._offsets[idx]
             return self._data[start:start + self._lengths[idx]]
 

nibabel/streamlines/tests/test_array_sequence.py

@@ -1,12 +1,12 @@
 import os
+import sys
 import unittest
 import tempfile
 import numpy as np
 
 from nose.tools import assert_equal, assert_raises, assert_true
 from nibabel.testing import assert_arrays_equal
 from numpy.testing import assert_array_equal
-from nibabel.externals.six.moves import zip, zip_longest
 
 from ..array_sequence import ArraySequence, is_array_sequence
 
@@ -32,7 +32,8 @@ def check_empty_arr_seq(seq):
     assert_equal(len(seq), 0)
     assert_equal(len(seq._offsets), 0)
     assert_equal(len(seq._lengths), 0)
-    assert_equal(seq._data.ndim, 0)
+    # assert_equal(seq._data.ndim, 0)
+    assert_equal(seq._data.ndim, 1)
     assert_true(seq.common_shape == ())
 
 
@@ -138,6 +139,11 @@ def test_arraysequence_append(self):
         seq.append(element)
         check_arr_seq(seq, SEQ_DATA['data'] + [element])
 
+        # Append a list of list.
+        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
+        seq.append(element.tolist())
+        check_arr_seq(seq, SEQ_DATA['data'] + [element])
+
         # Append to an empty ArraySequence.
         seq = ArraySequence()
         seq.append(element)
@@ -164,6 +170,11 @@ def test_arraysequence_extend(self):
         seq.extend(new_data)
         check_arr_seq(seq, SEQ_DATA['data'] + new_data)
 
+        # Extend with a generator.
+        seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
+        seq.extend((d for d in new_data))
+        check_arr_seq(seq, SEQ_DATA['data'] + new_data)
+
         # Extend with another `ArraySequence` object.
         seq = SEQ_DATA['seq'].copy()  # Copy because of in-place modification.
         seq.extend(ArraySequence(new_data))
@@ -195,6 +206,9 @@ def test_arraysequence_getitem(self):
         for i, e in enumerate(SEQ_DATA['seq']):
             assert_array_equal(SEQ_DATA['seq'][i], e)
 
+            if sys.version_info < (3,):
+                assert_array_equal(SEQ_DATA['seq'][long(i)], e)
+
         # Get all items using indexing (creates a view).
         indices = list(range(len(SEQ_DATA['seq'])))
         seq_view = SEQ_DATA['seq'][indices]

nibabel/streamlines/tests/test_streamlines.py

@@ -226,23 +226,28 @@ def test_save_tractogram_file(self):
             assert_true(issubclass(w[0].category, ExtensionWarning))
             assert_true("extension" in str(w[0].message))
 
+        with InTemporaryDirectory():
+            nib.streamlines.save(trk_file, "dummy.trk")
+            tfile = nib.streamlines.load("dummy.trk", lazy_load=False)
+            assert_tractogram_equal(tfile.tractogram, tractogram)
+
     def test_save_empty_file(self):
         tractogram = Tractogram()
         for ext, cls in nib.streamlines.FORMATS.items():
             with InTemporaryDirectory():
-                with open('streamlines' + ext, 'w+b') as f:
-                    nib.streamlines.save(tractogram, f.name)
-                    tfile = nib.streamlines.load(f, lazy_load=False)
-                    assert_tractogram_equal(tfile.tractogram, tractogram)
+                filename = 'streamlines' + ext
+                nib.streamlines.save(tractogram, filename)
+                tfile = nib.streamlines.load(filename, lazy_load=False)
+                assert_tractogram_equal(tfile.tractogram, tractogram)
 
     def test_save_simple_file(self):
         tractogram = Tractogram(DATA['streamlines'])
         for ext, cls in nib.streamlines.FORMATS.items():
             with InTemporaryDirectory():
-                with open('streamlines' + ext, 'w+b') as f:
-                    nib.streamlines.save(tractogram, f.name)
-                    tfile = nib.streamlines.load(f, lazy_load=False)
-                    assert_tractogram_equal(tfile.tractogram, tractogram)
+                filename = 'streamlines' + ext
+                nib.streamlines.save(tractogram, filename)
+                tfile = nib.streamlines.load(filename, lazy_load=False)
+                assert_tractogram_equal(tfile.tractogram, tractogram)
 
     def test_save_complex_file(self):
         complex_tractogram = Tractogram(DATA['streamlines'],
@@ -251,18 +256,19 @@ def test_save_complex_file(self):
 
         for ext, cls in nib.streamlines.FORMATS.items():
             with InTemporaryDirectory():
-                with open('streamlines' + ext, 'w+b') as f:
-                    with clear_and_catch_warnings(record=True,
-                                                  modules=[trk]) as w:
-                        nib.streamlines.save(complex_tractogram, f.name)
-
-                        # If streamlines format does not support saving data
-                        # per point or data per streamline, a warning message
-                        # should be issued.
-                        if not (cls.support_data_per_point() and
-                                cls.support_data_per_streamline()):
-                            assert_equal(len(w), 1)
-                            assert_true(issubclass(w[0].category, Warning))
+                filename = 'streamlines' + ext
+
+                with clear_and_catch_warnings(record=True,
+                                              modules=[trk]) as w:
+                    nib.streamlines.save(complex_tractogram, filename)
+
+                    # If streamlines format does not support saving data
+                    # per point or data per streamline, a warning message
+                    # should be issued.
+                    if not (cls.support_data_per_point() and
+                            cls.support_data_per_streamline()):
+                        assert_equal(len(w), 1)
+                        assert_true(issubclass(w[0].category, Warning))
 
                     tractogram = Tractogram(DATA['streamlines'])
 
@@ -272,7 +278,7 @@ def test_save_complex_file(self):
                     if cls.support_data_per_streamline():
                         tractogram.data_per_streamline = DATA['data_per_streamline']
 
-                    tfile = nib.streamlines.load(f, lazy_load=False)
+                    tfile = nib.streamlines.load(filename, lazy_load=False)
                     assert_tractogram_equal(tfile.tractogram, tractogram)
 
     def test_load_unknown_format(self):

nibabel/streamlines/tests/test_tractogram.py

@@ -1,3 +1,4 @@
+import sys
 import unittest
 import numpy as np
 import warnings
@@ -241,6 +242,22 @@ def test_tractogram_creation(self):
                          DATA['data_per_streamline'],
                          DATA['data_per_point'])
 
+        # Create a tractogram from another tractogram attributes.
+        tractogram2 = Tractogram(tractogram.streamlines,
+                                 tractogram.data_per_streamline,
+                                 tractogram.data_per_point)
+
+        assert_tractogram_equal(tractogram2, tractogram)
+
+        # Create a tractogram from a LazyTractogram object.
+        tractogram = LazyTractogram(DATA['streamlines_func'],
+                                    DATA['data_per_streamline_func'],
+                                    DATA['data_per_point_func'])
+
+        tractogram2 = Tractogram(tractogram.streamlines,
+                                 tractogram.data_per_streamline,
+                                 tractogram.data_per_point)
+
         # Inconsistent number of scalars between streamlines
         wrong_data = [[(1, 0, 0)]*1,
                       [(0, 1, 0), (0, 1)],
@@ -264,6 +281,9 @@ def test_tractogram_getitem(self):
         for i, t in enumerate(DATA['tractogram']):
             assert_tractogram_item_equal(DATA['tractogram'][i], t)
 
+            if sys.version_info < (3,):
+                assert_tractogram_item_equal(DATA['tractogram'][long(i)], t)
+
         # Get one TractogramItem out of two.
         tractogram_view = DATA['simple_tractogram'][::2]
         check_tractogram(tractogram_view, DATA['streamlines'][::2])
@@ -411,7 +431,7 @@ def test_lazy_tractogram_creation(self):
                                'mean_colors': (x for x in DATA['mean_colors'])}
 
         # Creating LazyTractogram with generators is not allowed as
-        # generators get exhausted and are not reusable unlike coroutines.
+        # generators get exhausted and are not reusable unlike generator function.
         assert_raises(TypeError, LazyTractogram, streamlines)
         assert_raises(TypeError, LazyTractogram,
                       data_per_streamline=data_per_streamline)
@@ -430,7 +450,7 @@ def test_lazy_tractogram_creation(self):
         assert_true(check_iteration(tractogram))
         assert_equal(len(tractogram), len(DATA['streamlines']))
 
-        # Coroutines get re-called and creates new iterators.
+        # Generator functions get re-called and creates new iterators.
         for i in range(2):
             assert_tractogram_equal(tractogram, DATA['tractogram'])
 
@@ -441,7 +461,7 @@ def test_lazy_tractogram_create_from(self):
         tractogram = LazyTractogram.create_from(_empty_data_gen)
         check_tractogram(tractogram)
 
-        # Create `LazyTractogram` from a coroutine yielding TractogramItem
+        # Create `LazyTractogram` from a generator function yielding TractogramItem.
         data = [DATA['streamlines'], DATA['fa'], DATA['colors'],
                 DATA['mean_curvature'], DATA['mean_torsion'],
                 DATA['mean_colors']]
@@ -530,13 +550,13 @@ def test_lazy_tractogram_copy(self):
         # Check we copied the data and not simply created new references.
         assert_true(tractogram is not DATA['lazy_tractogram'])
 
-        # When copying LazyTractogram, coroutines generating streamlines should
-        # be the same.
+        # When copying LazyTractogram, the generator function yielding streamlines
+        # should stay the same.
         assert_true(tractogram._streamlines
                     is DATA['lazy_tractogram']._streamlines)
 
         # Copying LazyTractogram, creates new internal LazyDict objects,
-        # but coroutines contained in it should be the same.
+        # but generator functions contained in it should stay the same.
         assert_true(tractogram._data_per_streamline
                     is not DATA['lazy_tractogram']._data_per_streamline)
         assert_true(tractogram._data_per_point