Merge pull request #41 from dwf/equizip

Added equizip + tests.

picklable_itertools/extras.py

@@ -8,6 +8,7 @@
 """
 import six
 from .base import BaseItertool
+from .map_zip import izip_longest
 from .iter_dispatch import iter_
 
 
@@ -74,3 +75,29 @@ def __next__(self):
         if len(items) == 0:
             raise StopIteration
         return tuple(items)
+
+
+class NoMoreItems(object):
+    """Sentinel value for `equizip`. Do not use for any other purpose."""
+    pass
+
+
+class IterableLengthMismatch(ValueError):
+    """Raised if an iterator passed to `equizip` is shorter than others."""
+    pass
+
+
+class equizip(izip_longest):
+    """Like `izip_longest` but ensures the sequences are the same length.
+
+    Raises :class:`IterableLengthMismatch` if one of the iterators
+    terminates prematurely.
+    """
+    def __init__(self, *args):
+        super(equizip, self).__init__(*args, fillvalue=NoMoreItems)
+
+    def __next__(self):
+        next_item = super(equizip, self).__next__()
+        if NoMoreItems in next_item:
+            raise IterableLengthMismatch
+        return next_item

tests/test_extras.py

@@ -1,5 +1,8 @@
 from unittest import SkipTest
-from picklable_itertools.extras import partition, partition_all
+from nose.tools import assert_raises
+from six.moves import zip
+from picklable_itertools.extras import (partition, partition_all,
+                                        IterableLengthMismatch, equizip)
 
 from . import verify_same, verify_pickle
 
@@ -17,3 +20,10 @@ def test_partition():
         yield verify_same, obj, ref, None, 3, [5, 9, 2, 9, 2]
         yield verify_same, obj, ref, None, 3, [5, 9, 2, 9, 2]
         yield verify_pickle, obj, ref, 2, 1, 3, [5, 9, 2, 9, 2, 4, 3]
+
+
+def test_equizip():
+    yield verify_same, equizip, zip, None, [3, 4], [9, 2], [9, 9]
+    yield verify_same, equizip, zip, None, [3, 4, 8, 4, 2]
+    assert_raises(IterableLengthMismatch, list, equizip([5, 4, 3], [2, 1]))
+    assert_raises(IterableLengthMismatch, list, equizip([5, 4, 3], []))