add slicing support to memo subscripting in memoized generators

Technologicat · Technologicat · commit 7fa7c98abd55 · 2021-06-16T11:10:20.000+03:00
diff --git a/unpythonic/gmemo.py b/unpythonic/gmemo.py
@@ -132,12 +132,24 @@ def __next__(self):
             if kind is _fail:
                 raise value
             return value
-    # Support the `collections.abc.Sequence` API for already-computed items
+    # Support a subset of the `collections.abc.Sequence` API for already-computed items
     def __len__(self):
         return len(self.memo)
     def __getitem__(self, k):
-        if k >= len(self.memo):
-            raise IndexError(f"Attempted to access index {k} of memoized generator; only {len(self.memo)} items available (at least so far)")
+        if not isinstance(k, (int, slice)):
+            raise TypeError(f"Expected an int or slice index, got {type(k)} with value {repr(k)}")
+        length = len(self.memo)
+        if isinstance(k, slice):
+            # For slices where at least one item raises an exception, we raise the
+            # exception that is encountered first when walking the slice.
+            lst = []
+            for kind, value in self.memo[k]:
+                if kind is _fail:
+                    raise value
+                lst.append(value)
+            return lst
+        if k >= length or k < -length:
+            raise IndexError(f"memoized generator index out of range; got {k}, with {len(self.memo)} items currently available")
         kind, value = self.memo[k]
         if kind is _fail:
             raise value
diff --git a/unpythonic/tests/test_gmemo.py b/unpythonic/tests/test_gmemo.py
@@ -97,6 +97,9 @@ def gen():
         def gen():
             yield from range(5)
         g3 = gen()
+
+        # Any item that has entered the memo can be retrieved by subscripting.
+        # len() is the current length of the memo.
         test[len(g3) == 0]
         next(g3)
         test[len(g3) == 1]
@@ -107,8 +110,29 @@ def gen():
         test[g3[0] == 0]
         test[g3[1] == 1]
         test[g3[2] == 2]
+
+        # Items not yet memoized cannot be retrieved from the memo.
         test_raises[IndexError, g3[3]]
 
+        # Negative indices work too, counting from the current end of the memo.
+        test[g3[-1] == 2]
+        test[g3[-2] == 1]
+        test[g3[-3] == 0]
+
+        # Counting back past the start is an error, just like in `list`.
+        test_raises[IndexError, g3[-4]]
+
+        # Slicing is supported.
+        test[g3[0:3] == [0, 1, 2]]
+        test[g3[0:2] == [0, 1]]
+        test[g3[::-1] == [2, 1, 0]]
+        test[g3[0::2] == [0, 2]]
+        test[g3[2::-2] == [2, 0]]
+
+        # Out-of-range slices produce the empty list, like in `list`.
+        test[g3[3:] == []]
+        test[g3[-4::-1] == []]
+
     with testset("memoizing a sequence partially"):
         # To do this, build a chain of generators, then memoize only the last one:
         evaluations = Counter()
