Add take / takeWhile

README.md

@@ -450,10 +450,10 @@ R.subtract(date(1,2,3), date(1,2,3)) # float('nan)
 - [ ] symmetricDifferenceWith
 - [x] T
 - [x] tail
-- [ ] take
+- [x] take
 - [ ] takeLast
 - [ ] takeLastWhile
-- [ ] takeWhile
+- [x] takeWhile
 - [ ] tap
 - [ ] test
 - [ ] thunkify

ramda/__init__.py

@@ -78,6 +78,8 @@
 from .sum import sum
 from .T import T
 from .tail import tail
+from .take import take
+from .takeWhile import takeWhile
 from .toString import toString
 from .uniq import uniq
 from .uniqBy import uniqBy

ramda/private/_xtake.py

@@ -0,0 +1,21 @@
+from ._helper import getAttribute
+from ._reduced import _reduced
+from ._xfBase import XfBase
+
+
+class XTake(XfBase):
+  def __init__(self, n, xf):
+    self.xf = xf
+    self.n = n
+    self.i = 0
+
+  def step(self, result, input):
+    self.i += 1
+    ret = result if self.n == 0 else getAttribute(self.xf, '@@transducer/step')(result, input)
+    if self.n >= 0 and self.i >= self.n:
+      return _reduced(ret)
+    else:
+      return ret
+
+
+def _xtake(n): return lambda xf: XTake(n, xf)

ramda/private/_xtakeWhile.py

@@ -0,0 +1,18 @@
+from ._helper import getAttribute
+from ._reduced import _reduced
+from ._xfBase import XfBase
+
+
+class XTakeWhile(XfBase):
+  def __init__(self, f, xf):
+    self.xf = xf
+    self.f = f
+
+  def step(self, result, input):
+    if self.f(input):
+      return getAttribute(self.xf, '@@transducer/step')(result, input)
+    else:
+      return _reduced(result)
+
+
+def _xtakeWhile(f): return lambda xf: XTakeWhile(f, xf)

ramda/take.py

@@ -0,0 +1,11 @@
+from .private._curry2 import _curry2
+from .private._dispatchable import _dispatchable
+from .private._xtake import _xtake
+from .slice import slice
+
+
+def inner_take(n, xs):
+  return slice(0, None if n < 0 else n, xs)
+
+
+take = _curry2(_dispatchable(['take'], _xtake, inner_take))

ramda/takeWhile.py

@@ -0,0 +1,15 @@
+from .private._curry2 import _curry2
+from .private._dispatchable import _dispatchable
+from .private._xtakeWhile import _xtakeWhile
+from .slice import slice
+
+
+def inner_takeWhile(fn, xs):
+  idx = 0
+  length = len(xs)
+  while idx < length and fn(xs[idx]):
+    idx += 1
+  return slice(0, idx, xs)
+
+
+takeWhile = _curry2(_dispatchable(['takeWhile'], _xtakeWhile, inner_takeWhile))

test/test_take.py

@@ -0,0 +1,59 @@
+
+import unittest
+from itertools import count
+
+import ramda as R
+
+"""
+https://github.com/ramda/ramda/blob/master/test/take.js
+"""
+
+
+class TestTake(unittest.TestCase):
+  def test_takes_only_the_first_n_elements_from_a_list(self):
+    self.assertEqual(['a', 'b', 'c'], R.take(3, ['a', 'b', 'c', 'd', 'e', 'f', 'g']))
+
+  def test_returns_only_as_many_as_the_array_can_provide(self):
+    self.assertEqual([1, 2], R.take(3, [1, 2]))
+    self.assertEqual([], R.take(3, []))
+
+  def test_returns_an_equivalent_list_if_n_is_negative(self):
+    self.assertEqual([1, 2, 3], R.take(-1, [1, 2, 3]))
+
+  def test_never_returns_the_input_array(self):
+    xs = [1, 2, 3]
+    self.assertIsNot(xs, R.take(3, xs))
+    self.assertIsNot(xs, R.take(-1, xs))
+
+  def test_can_operate_on_strings(self):
+    self.assertEqual('Ram', R.take(3, 'Ramda'))
+    self.assertEqual('Ra', R.take(2, 'Ramda'))
+    self.assertEqual('R', R.take(1, 'Ramda'))
+    self.assertEqual('', R.take(0, 'Ramda'))
+
+  def test_handles_zero_correctly(self):
+    self.assertEqual([], R.into([], R.take(0), [1, 2, 3]))
+
+  def test_steps_correct_number_of_times(self):
+    count = 0
+
+    def spy(_):
+      nonlocal count
+      count += 1
+
+    R.into([], R.compose(R.map(spy), R.take(2)), [1, 2, 3])
+    self.assertEqual(2, count)
+
+  def test_transducer_called_for_every_member_of_list_if_n_is_negative(self):
+    count = 0
+
+    def spy(_):
+      nonlocal count
+      count += 1
+
+    R.into([], R.compose(R.map(spy), R.take(-1)), [1, 2, 3])
+    self.assertEqual(3, count)
+
+
+if __name__ == '__main__':
+  unittest.main()

test/test_takeWhile.py

@@ -0,0 +1,31 @@
+
+import unittest
+
+import ramda as R
+
+"""
+https://github.com/ramda/ramda/blob/master/test/takeWhile.js
+"""
+
+
+class TestTakeWhile(unittest.TestCase):
+  def test_continues_taking_elements_while_the_function_reports_true(self):
+    self.assertEqual([1, 3], R.takeWhile(lambda x: x != 5, [1, 3, 5, 7, 9]))
+
+  def test_starts_at_the_right_arg_and_acknowledges_undefined(self):
+    self.assertEqual([], R.takeWhile(lambda: False, []))
+    self.assertEqual([1, 3], R.takeWhile(lambda x: x is not None, [1, 3, None, 5, 7]))
+
+  def test_can_operate_on_strings(self):
+    self.assertEqual('Ram', R.takeWhile(lambda x: x != 'd', 'Ramda'))
+
+  def test_can_act_as_a_transducer(self):
+    def isNotFour(x): return x != 4
+    input = [1, 2, 3, 4, 3, 2, 1]
+    expected = [1, 2, 3]
+    self.assertEqual(expected, R.into([], R.takeWhile(isNotFour), input))
+    # TODO: transducer
+
+
+if __name__ == '__main__':
+  unittest.main()