Merge pull request #399 from groutr/fix-partition-all

Avoid using .index() in partition_all

Author: eriknw

toolz/itertoolz.py

@@ -720,7 +720,23 @@ def partition_all(n, seq):
         yield prev
         prev = item
     if prev[-1] is no_pad:
-        yield prev[:prev.index(no_pad)]
+        try:
+            # If seq defines __len__, then
+            # we can quickly calculate where no_pad starts
+            yield prev[:len(seq) % n]
+        except TypeError:
+            # Get first index of no_pad without using .index()
+            # https://github.com/pytoolz/toolz/issues/387
+            # Binary search from CPython's bisect module,
+            # modified for identity testing.
+            lo, hi = 0, n
+            while lo < hi:
+                mid = (lo + hi) // 2
+                if prev[mid] is no_pad:
+                    hi = mid
+                else:
+                    lo = mid + 1
+            yield prev[:lo]
     else:
         yield prev
 

toolz/tests/test_itertoolz.py

@@ -318,6 +318,17 @@ def test_partition_all():
     assert list(partition_all(3, range(5))) == [(0, 1, 2), (3, 4)]
     assert list(partition_all(2, [])) == []
 
+    # Regression test: https://github.com/pytoolz/toolz/issues/387
+    class NoCompare(object):
+        def __eq__(self, other):
+            if self.__class__ == other.__class__:
+                return True
+            raise ValueError()
+    obj = NoCompare()
+    result = [(obj, obj, obj, obj), (obj, obj, obj)]
+    assert list(partition_all(4, [obj]*7)) == result
+    assert list(partition_all(4, iter([obj]*7))) == result
+
 
 def test_count():
     assert count((1, 2, 3)) == 3