+drop and take

orsinium · orsinium · commit dd8a71f44d5d · 2019-09-05T11:16:28.000+02:00
diff --git a/lambdas/_list.py b/lambdas/_list.py
@@ -1,7 +1,7 @@
-from ._bool import TRUE, FALSE
+from ._bool import TRUE, FALSE, OR
 from ._pair import CONS, CAR, CDR
 from ._combinators import Y
-from ._natural import INC, GTE
+from ._natural import DEC, INC, GTE, ISZERO
 
 
 LIST = CONS(TRUE)(TRUE)
@@ -53,3 +53,13 @@
     (l)
     (LIST)
 )
+DROP = lambda n: lambda l: n(TAIL)(l)
+TAKE = Y(lambda f: lambda n: lambda l: (
+    OR(EMPTY(l))(ISZERO(n))
+    (lambda _: LIST)
+    (lambda _: (
+        PREPEND(f(DEC(n))(TAIL(l)))
+        (HEAD(l))
+    ))
+    (TRUE)
+))
diff --git a/tests/test_list.py b/tests/test_list.py
@@ -2,7 +2,8 @@
 
 from lambdas import ZERO, ONE, TWO, THREE, FOUR
 from lambdas import ADD, LTE
-from lambdas import LIST, APPEND, PREPEND, REVERSE, MAP, RANGE, REDUCE, FILTER
+from lambdas import LIST, APPEND, PREPEND, REVERSE
+from lambdas import MAP, RANGE, REDUCE, FILTER, TAKE, DROP
 from lambdas import decode_number, decode_list
 
 
@@ -97,3 +98,28 @@ def test_filter(given, expected):
     result = FILTER(LTE(TWO))(lst)
     decoded = [decode_number(n) for n in decode_list(result)]
     assert decoded == expected
+
+
+@pytest.mark.parametrize('given, number, expected', [
+    ([1, 2, 3], ONE, [2, 3]),
+])
+def test_drop(given, number, expected):
+    lst = LIST
+    for el in given:
+        lst = APPEND(lst)(el)
+    result = DROP(number)(lst)
+    assert decode_list(result) == expected
+
+
+@pytest.mark.parametrize('given, number, expected', [
+    ([1, 2, 3], TWO, [1, 2]),
+    ([1, 2, 3], FOUR, [1, 2, 3]),
+    ([1, 2, 3], ZERO, []),
+    ([], TWO, []),
+])
+def test_take(given, number, expected):
+    lst = LIST
+    for el in given:
+        lst = APPEND(lst)(el)
+    result = TAKE(number)(lst)
+    assert decode_list(result) == expected
diff --git a/tests/test_natural.py b/tests/test_natural.py
@@ -1,7 +1,7 @@
 import pytest
 
 from lambdas import TRUE, FALSE
-from lambdas import ZERO, ONE, TWO, THREE, FOUR, FIVE, SIX
+from lambdas import ZERO, ONE, TWO, THREE, FOUR, FIVE, SIX, SEVEN, EIGHT, NINE, TEN
 from lambdas import ADD, INC, MUL, POW, DEC, SUB
 from lambdas import ISZERO, GTE, LTE, GT, LT, EQ
 from lambdas import decode_number
@@ -14,6 +14,10 @@
     (FOUR,  4),
     (FIVE,  5),
     (SIX,   6),
+    (SEVEN, 7),
+    (EIGHT, 8),
+    (NINE,  9),
+    (TEN,   10),
 ])
 def test_numbers(given, expected):
     assert decode_number(given) == expected
