Merge pull request #7 from brandonwillard/use-custom-getitem

Make cdr use custom __getitem__ methods

Author: brandonwillard

cons/core.py

@@ -14,10 +14,12 @@
 
 
 def rest(seq):
-    if isinstance(seq, Iterator) and length_hint(seq, 2) <= 1:
+    if isinstance(seq, Sequence):
+        return seq[1:]
+    elif isinstance(seq, Iterator) and length_hint(seq, 2) <= 1:
         return iter([])
-    else:
-        return islice(seq, 1, None)
+
+    return islice(seq, 1, None)
 
 
 class ConsError(ValueError):

tests/test_cons.py

@@ -209,6 +209,15 @@ def test_car_cdr():
     assert car(cons(1, cons("a", "b"))) == 1
     assert cdr(cons(1, cons("a", "b"))) == cons("a", "b")
 
+    # We need to make sure that `__getitem__` is actually used.
+    from collections import UserList
+
+    class CustomList(UserList):
+        def __getitem__(self, *args):
+            return [5]
+
+    assert cdr(CustomList([1, 2, 3])) == [5]
+
 
 def test_unification():
     car_lv, cdr_lv = var(), var()