Add negative start, stop support for unpythonic.slicing.islice

Technologicat · Technologicat · commit 9ab1afa9054c · 2019-08-29T00:02:57.000+03:00
Resolves #25.
diff --git a/unpythonic/slicing.py b/unpythonic/slicing.py
@@ -6,20 +6,28 @@
 from itertools import islice as islicef
 
 from .fup import fupdate
+from .it import first, lastn, butlastn
+from .misc import CountingIterator
 
 def islice(iterable):
-    """Use itertools.islice with slice syntax.
+    """Use itertools.islice with slice syntax, with some bonus features.
 
     Usage::
 
         islice(iterable)[idx_or_slice]
 
-    The slicing variant calls ``itertools.islice`` with the corresponding
-    slicing parameters.
+    For convenience:
 
-    As a convenience feature: a single index is interpreted as a length-1 islice
-    starting at that index. The slice is then immediately evaluated and the item
-    is returned.
+      - Negative ``start``, ``stop`` are supported. **CAUTION**: using a negative
+        start or stop will force the iterable, because that is the only way to
+        know its length.
+
+      - A single index (negative also allowed) is interpreted as a length-1
+        islice starting at that index. The slice is then immediately evaluated
+        and the item is returned.
+
+    Once negative indices have been handled, the slicing variant calls
+    ``itertools.islice`` with the corresponding slicing parameters.
 
     Examples::
 
@@ -37,9 +45,10 @@ def islice(iterable):
         assert tuple(islice(odds)[:5]) == (1, 3, 5, 7, 9)
         assert tuple(islice(odds)[:5]) == (11, 13, 15, 17, 19)  # five more
 
-    **CAUTION**: Keep in mind ``itertools.islice`` does not support negative
-    indexing for any of ``start``, ``stop`` or ``step``, and that the slicing
-    process consumes elements from the iterable.
+    **CAUTION**: Keep in mind the slicing process consumes elements from the
+    iterable.
+
+    **CAUTION**: ``step``, if present, must be positive.
     """
     # manually curry to take indices later, but expect them in subscript syntax to support slicing
     class islice1:
@@ -48,10 +57,50 @@ def __getitem__(self, k):
             if isinstance(k, tuple):
                 raise TypeError("multidimensional indexing not supported, got {}".format(k))
             if isinstance(k, slice):
-                return islicef(iterable, k.start, k.stop, k.step)
-            return tuple(islicef(iterable, k, k + 1))[0]
+                start, stop, step = k.start, k.stop, k.step
+                it = iter(iterable)
+                # One or both of start and stop may be negative or None.
+                # Step must be positive; filter first, then slice normally.
+                #
+                # A general iterable doesn't know its length (might not even be
+                # knowable; it may be a generator), so if we get a negative
+                # start or stop, the only way to find the correct position is
+                # to force the iterable until it ends (if ever).
+                if start and start < 0 and stop and stop < 0:
+                    it = butlastn(-stop, lastn(-start, iterable))
+                    start = stop = None
+                elif start and start < 0:
+                    n, start = -start, None
+                    if not stop:
+                        it = lastn(n, iterable)
+                    else:  # stop and stop > 0:
+                        # to adjust stop, we must know how many items are dropped
+                        cit = CountingIterator(iterable)
+                        it = tuple(lastn(n, cit))  # force to actually count (note lastn stores only <= n items)
+                        n_dropped = max(0, cit.count - n)  # max needed if start is past the start of iterable
+                        stop -= n_dropped
+                        assert stop >= 0
+                elif stop and stop < 0:
+                    it = butlastn(-stop, iterable)
+                    stop = None
+                return islicef(it, start, stop, step)
+            if k < 0:
+                return first(lastn(-k, iterable))
+            return first(islicef(iterable, k, k + 1))
     return islice1()
 
+# Basic idea, no negative index support:
+# def islice(iterable):
+#     class islice1:
+#         """Subscript me to perform the slicing."""
+#         def __getitem__(self, k):
+#             if isinstance(k, tuple):
+#                 raise TypeError("multidimensional indexing not supported, got {}".format(k))
+#             if isinstance(k, slice):
+#                 return islicef(iterable, k.start, k.stop, k.step)
+#             return first(islicef(iterable, k, k + 1))
+#     return islice1()
+
 def fup(seq):
     """Functionally update a sequence.
 
diff --git a/unpythonic/test/test_slicing.py b/unpythonic/test/test_slicing.py
@@ -30,6 +30,34 @@ def test():
     assert tuple(islice(odds)[:5]) == (1, 3, 5, 7, 9)
     assert tuple(islice(odds)[:5]) == (11, 13, 15, 17, 19)  # five more
 
+    # negative start, stop in islice
+    #
+    # !! step must be positive !!
+    #
+    # CAUTION: will force the iterable (at the latest at the time when the
+    # first item is read from it), since that's the only way to know where it
+    # ends, if at all.
+    assert tuple(islice(range(10))[-3:]) == (7, 8, 9)       # start < 0, no stop
+    assert tuple(islice(range(10))[-3:9]) == (7, 8)         # start < 0, stop > 0, before end
+    assert tuple(islice(range(10))[-2:10]) == (8, 9)        # start < 0, stop > 0, at end
+    assert tuple(islice(range(10))[-2:20]) == (8, 9)        # start < 0, stop > 0, beyond end
+
+    assert tuple(islice(range(10))[:-8]) == (0, 1)          # no start, stop < 0
+    assert tuple(islice(range(10))[6:-2]) == (6, 7)         # start > 0, stop < 0
+    assert tuple(islice(range(10))[10:-2]) == ()            # start > 0, at end, stop < 0
+    assert tuple(islice(range(10))[20:-2]) == ()            # start > 0, beyond end, stop < 0
+
+    assert tuple(islice(range(10))[-8:-4]) == (2, 3, 4, 5)  # start < 0, stop < 0
+
+    assert tuple(islice(range(10))[-6::2]) == (4, 6, 8)     # step
+    assert tuple(islice(range(10))[:-2:2]) == (0, 2, 4, 6)  # step
+    assert tuple(islice(range(10))[-8:-2:3]) == (2, 5)      # step
+
+    # edge cases, should behave like list does
+    assert tuple(islice(range(10))[:-20]) == ()                # stop < 0, past the start of the iterable
+    assert tuple(islice(range(10))[-20:]) == tuple(range(10))  # start < 0, past the start of the iterable
+    assert tuple(islice(range(10))[-20:5]) == tuple(range(5))  # same, but with stop > 0
+
     print("All tests PASSED")
 
 if __name__ == '__main__':
