RF: created peek_next util function

Author: MarcCote

nibabel/streamlines/tck.py

@@ -20,6 +20,7 @@
 from .tractogram_file import HeaderError, DataError
 from .tractogram import TractogramItem, Tractogram, LazyTractogram
 from .header import Field
+from .utils import peek_next
 
 MEGABYTE = 1024 * 1024
 
@@ -201,9 +202,7 @@ def save(self, fileobj):
                 # Use the first element to check
                 #  1) the tractogram is not empty;
                 #  2) quantity of information saved along each streamline.
-                first_item = next(tractogram)
-                # Put back the first element at its place.
-                tractogram = itertools.chain([first_item], tractogram)
+                first_item, tractogram = peek_next(tractogram)
             except StopIteration:
                 # Empty tractogram
                 header[Field.NB_STREAMLINES] = 0

nibabel/streamlines/trk.py

@@ -22,6 +22,7 @@
 from .tractogram_file import DataError, HeaderError, HeaderWarning
 from .tractogram import TractogramItem, Tractogram, LazyTractogram
 from .header import Field
+from .utils import peek_next
 
 
 MAX_NB_NAMED_SCALARS_PER_POINT = 10
@@ -435,9 +436,7 @@ def save(self, fileobj):
                 # Use the first element to check
                 #  1) the tractogram is not empty;
                 #  2) quantity of information saved along each streamline.
-                first_item = next(tractogram)
-                # Put back the first element at its place.
-                tractogram = itertools.chain([first_item], tractogram)
+                first_item, tractogram = peek_next(tractogram)
             except StopIteration:
                 # Empty tractogram
                 header[Field.NB_STREAMLINES] = 0

nibabel/streamlines/utils.py

@@ -1,3 +1,5 @@
+import itertools
+
 import nibabel
 
 
@@ -29,3 +31,22 @@ def get_affine_from_reference(ref):
 
     # Assume `ref` is the name of a neuroimaging file.
     return nibabel.load(ref).affine
+
+
+def peek_next(iterable):
+    """ Peek next element of iterable.
+
+    Parameters
+    ----------
+    iterable
+        Iterable to peek the next element from.
+
+    Returns
+    -------
+    next_item
+        Element peeked from `iterable`.
+    new_iterable
+        Iterable behaving like if the original `iterable` was untouched.
+    """
+    next_item = next(iterable)
+    return next_item, itertools.chain([next_item], iterable)